#!/usr/bin/env python3
"""增量抓取九年义务教育全科目中文语料（维基百科 API）。

特性：
  - 增量：已有词条跳过，只补缺失的
  - 429 自动重试：等待 10 秒后重试，最多 5 次
  - 长延时（3秒）防限流
  - 每词条截取前 3000 字符
  - 断点续传：Ctrl+C 中断后再次运行会跳过已抓的

用法：
  python scripts/fetch_corpus_incremental.py           # 抓全部科目
  python scripts/fetch_corpus_incremental.py 语文 数学   # 只抓指定科目
"""
import json
import urllib.request
import urllib.parse
import time
import re
import sys
from pathlib import Path

CORPUS_DIR = Path(__file__).parent.parent / 'corpus'
CORPUS_DIR.mkdir(parents=True, exist_ok=True)

# 九年义务教育各科目核心词条（覆盖小学到初中知识点）
SUBJECTS = {
    '语文': [
        # 古诗词与文言文（部编版教材重点篇目）
        '静夜思', '春晓', '登鹳雀楼', '望庐山瀑布', '绝句', '春望',
        '江雪', '游子吟', '清明', '枫桥夜泊', '寻隐者不遇',
        '咏鹅', '回乡偶书', '咏柳', '凉州词', '出塞',
        '芙蓉楼送辛渐', '送元二使安西', '九月九日忆山东兄弟',
        '赠汪伦', '黄鹤楼送孟浩然之广陵', '望天门山',
        '别董大', '早发白帝城', '赠花卿', '江南逢李龟年',
        '春夜喜雨', '江畔独步寻花', '渔歌子', '塞下曲',
        '望洞庭', '浪淘沙', '赋得古原草送别', '忆江南',
        '悯农', '蚕妇', '陶者', '梅花', '元日',
        '泊船瓜洲', '书湖阴先生壁', '六月二十七日望湖楼醉书',
        '饮湖上初晴后雨', '题西林壁', '惠崇春江晚景',
        '晓出净慈寺送林子方', '小池', '宿新市徐公店',
        '春日', '观书有感', '题临安邸', '游园不值',
        '乡村四月', '墨梅', '石灰吟', '竹石', '所见',
        '村居', '己亥杂诗',
        # 文言文
        '论语', '孟子', '大学', '中庸', '诗经',
        '世说新语', '聊斋志异', '西游记', '水浒传',
        '三国演义', '红楼梦',
        '岳阳楼记', '醉翁亭记', '爱莲说', '陋室铭',
        '桃花源记', '小石潭记', '记承天寺夜游',
        '送东阳马生序', '曹刿论战', '邹忌讽齐王纳谏',
        '出师表', '陈涉世家',
        # 作家
        '李白', '杜甫', '白居易', '王维', '苏轼',
        '李清照', '辛弃疾', '陆游', '鲁迅', '朱自清',
        '老舍', '巴金', '冰心',
        # 文学常识
        '唐诗', '宋词', '元曲', '文言文', '成语',
    ],
    '数学': [
        '加法', '减法', '乘法', '除法', '四则运算',
        '分数', '小数', '百分比', '比例', '方程',
        '一元一次方程', '二元一次方程', '一元二次方程',
        '不等式', '一元一次不等式',
        '几何学', '三角形', '圆', '正方形', '长方形',
        '平行四边形', '梯形', '多边形', '正方体', '长方体',
        '圆柱', '圆锥', '球', '扇形',
        '面积', '体积', '周长', '角度', '弧',
        '整数', '有理数', '无理数', '实数', '复数',
        '代数', '函数', '一次函数', '二次函数', '反比例函数',
        '勾股定理', '圆周率', '素数', '因数', '倍数',
        '最大公约数', '最小公倍数', '绝对值', '平方根',
        '统计', '概率', '平均数', '中位数', '众数',
        '负数', '相反数', '倒数', '科学记数法',
        '相似三角形', '全等三角形', '平行线',
        '轴对称', '中心对称', '平移', '旋转',
        '坐标', '直角坐标系', '象限',
    ],
    '英语': [
        '英文字母', '英语', '英语语法', '英语动词',
        '英语名词', '英语形容词', '英语副词', '英语代词',
        '英语介词', '英语连词', '英语冠词',
        '英语时态', '一般现在时', '现在进行时',
        '一般过去时', '过去进行时', '现在完成时',
        '一般将来时', '过去完成时', '过去完成进行时',
        '英语被动语态', '英语从句', '定语从句',
        '名词性从句', '状语从句', '宾语从句',
        '主语从句', '表语从句', '同位语从句',
        '英语虚拟语气', '英语不定式', '英语动名词', '英语分词',
        '英语单词', '英语发音', '英语拼写', '英语音标',
    ],
    '化学': [
        '化学', '化学元素', '元素周期表', '原子', '分子',
        '离子', '化合物', '混合物', '纯净物', '单质',
        '氧气', '氢气', '氮气', '二氧化碳', '一氧化碳',
        '水', '盐酸', '硫酸', '硝酸', '氢氧化钠',
        '碳酸钙', '氯化钠', '碳酸氢钠', '氢氧化钙',
        '酸', '碱', '盐', '氧化物', '过氧化物',
        '化学反应', '化合反应', '分解反应', '置换反应',
        '氧化还原反应', '复分解反应', '燃烧', '爆炸',
        '溶液', '溶质', '溶剂', '溶解度', '浓度',
        '金属', '合金', '铁', '铜', '铝', '锌', '银',
        '碳', '金刚石', '石墨', '有机化合物', '甲烷',
        '乙醇', '乙酸', '蛋白质', '糖类', '脂肪',
        '化学式', '化学方程式', '摩尔', '相对原子质量',
        '原子结构', '电子', '质子', '中子', '同位素',
        '化学键', '离子键', '共价键', '金属键',
    ],
    '政治': [
        '公民', '权利', '义务', '法律', '宪法',
        '刑法', '民法', '消费者权益', '未成年人保护',
        '人民代表大会制度', '全国人民代表大会',
        '中国共产党', '社会主义', '社会主义市场经济',
        '改革开放', '社会主义核心价值观',
        '国家安全', '爱国主义', '集体主义',
        '民主', '自由', '平等', '公正', '法治',
        '民族团结', '一国两制', '和平共处五项原则',
        '联合国', '世界贸易组织', '经济全球化',
        '人民代表大会', '政治协商制度', '民族区域自治',
        '基层群众自治', '依法治国', '以德治国',
    ],
    '历史': [
        '中国历史', '夏朝', '商朝', '西周', '东周',
        '春秋时期', '战国时期', '秦朝', '西汉', '东汉',
        '三国', '西晋', '东晋', '南北朝', '隋朝', '唐朝',
        '宋朝', '辽朝', '金朝', '元朝', '明朝', '清朝',
        '中华民国', '中华人民共和国',
        '鸦片战争', '太平天国运动', '洋务运动', '甲午战争',
        '戊戌变法', '义和团运动', '辛亥革命',
        '新文化运动', '五四运动', '中国共产党成立',
        '北伐战争', '土地革命', '长征',
        '抗日战争', '解放战争', '新中国成立',
        '丝绸之路', '四大发明', '郑和下西洋',
        '孔子', '老子', '孟子', '庄子', '韩非子', '墨子',
        '秦始皇', '汉武帝', '唐太宗', '宋太祖',
        '成吉思汗', '朱元璋', '康熙帝',
        '世界历史', '古埃及', '古希腊', '古罗马',
        '文艺复兴', '工业革命', '第一次世界大战',
        '第二次世界大战', '冷战', '欧洲联盟',
    ],
    '地理': [
        '地球', '地球仪', '经线', '纬线', '赤道',
        '本初子午线', '时区', '国际日期变更线',
        '大气层', '对流层', '水循环', '气候',
        '热带', '温带', '寒带', '季风', '降水',
        '地形', '平原', '高原', '山地', '盆地', '丘陵',
        '河流', '湖泊', '海洋', '海峡', '岛屿', '半岛',
        '中国地理', '中国行政区划', '省级行政区',
        '长江', '黄河', '珠江', '淮河', '黑龙江',
        '青藏高原', '内蒙古高原', '黄土高原', '云贵高原',
        '塔里木盆地', '准噶尔盆地', '柴达木盆地', '四川盆地',
        '华北平原', '东北平原', '长江中下游平原',
        '喜马拉雅山脉', '昆仑山脉', '秦岭',
        '北京', '上海', '广州', '深圳', '天津', '重庆',
        '香港', '澳门', '台湾',
        '世界地理', '亚洲', '欧洲', '非洲',
        '北美洲', '南美洲', '大洋洲', '南极洲',
        '太平洋', '大西洋', '印度洋', '北冰洋',
        '撒哈拉沙漠', '亚马逊雨林', '尼罗河',
        '人口', '城市', '农业', '工业', '交通运输',
    ],
    '物理': [
        '物理学', '力', '运动', '速度', '加速度',
        '牛顿运动定律', '牛顿第一定律', '牛顿第二定律',
        '牛顿第三定律', '惯性', '质量', '密度',
        '重力', '摩擦力', '弹力', '浮力', '压强',
        '大气压强', '液体压强', '液压千斤顶',
        '功', '功率', '能量', '动能', '势能',
        '机械能', '机械能守恒', '能量守恒定律',
        '杠杆', '滑轮', '斜面', '轮轴',
        '电', '电荷', '电流', '电压', '电阻',
        '欧姆定律', '串联', '并联', '电功率', '电能',
        '磁', '磁场', '磁体', '电磁感应', '电磁铁',
        '电动机', '发电机', '变压器',
        '光', '光的反射', '光的折射', '透镜',
        '凸透镜', '凹透镜', '光的色散', '颜色', '光谱',
        '声音', '声波', '音调', '响度', '音色',
        '回声', '超声波', '次声波', '噪声',
        '热', '温度', '温度计', '熔化', '凝固',
        '汽化', '液化', '升华', '凝华',
        '比热容', '热传导', '热对流', '热辐射',
        '物态变化', '沸腾', '蒸发', '内能',
        '比热', '热量', '热机', '热机效率',
    ],
    '生物': [
        '生物学', '细胞', '细胞膜', '细胞核', '细胞质',
        '细胞壁', '叶绿体', '线粒体', '染色体', 'DNA',
        '基因', '遗传', '变异', '遗传学',
        '光合作用', '呼吸作用', '新陈代谢',
        '植物', '动物', '细菌', '病毒', '真菌',
        '苔藓植物', '蕨类植物', '种子植物', '裸子植物', '被子植物',
        '脊椎动物', '无脊椎动物', '哺乳动物', '鸟类',
        '爬行动物', '两栖动物', '鱼类', '昆虫',
        '生态系统', '食物链', '食物网', '生物多样性',
        '生产者', '消费者', '分解者', '生物圈',
        '人体', '心脏', '血管', '血液', '血液循环',
        '肺', '呼吸', '肝脏', '胃', '肠',
        '神经系统', '大脑', '脊髓', '神经元',
        '消化系统', '吸收', '排泄', '泌尿系统',
        '生殖', '发育', '青春期', '激素',
        '群落', '种群', '生态系统稳定性',
        '进化', '自然选择', '达尔文', '物种起源',
        '组织', '器官', '系统', '细胞分裂', '细胞分化',
    ],
}


def fetch_extract(title: str, retries: int = 5) -> str:
    """从维基百科 API 获取词条纯文本摘要，429 自动重试"""
    api = "https://zh.wikipedia.org/w/api.php"
    params = urllib.parse.urlencode({
        'action': 'query',
        'titles': title,
        'prop': 'extracts',
        'explaintext': '1',
        'exsectionformat': 'plain',
        'format': 'json',
        'redirects': '1',
    })
    url = f"{api}?{params}"
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url, headers={
                'User-Agent': 'JspaceAI/1.0 (educational corpus fetcher; contact@example.com)'
            })
            with urllib.request.urlopen(req, timeout=20) as resp:
                data = json.loads(resp.read().decode('utf-8'))
            pages = data.get('query', {}).get('pages', {})
            for pid, page in pages.items():
                if pid != '-1' and 'extract' in page:
                    return page['extract'][:3000]  # 截取前 3000 字符
            return ''  # 词条不存在
        except urllib.error.HTTPError as e:
            if e.code == 429:
                wait = 10 * (attempt + 1)
                print(f"    429 限流，等待 {wait}s 后重试 ({attempt+1}/{retries})")
                time.sleep(wait)
                continue
            else:
                print(f"    HTTP {e.code}: {title}")
                return ''
        except Exception as e:
            print(f"    错误: {e}")
            time.sleep(5)
    return ''


def get_existing_titles(filepath: Path) -> set:
    """从已有文件读取已抓取的词条名"""
    if not filepath.exists():
        return set()
    content = filepath.read_text(encoding='utf-8')
    return set(re.findall(r'^## (.+)$', content, re.M))


def main():
    subjects = sys.argv[1:] if len(sys.argv) > 1 else list(SUBJECTS.keys())
    
    total_new = 0
    total_skip = 0
    
    for subject in subjects:
        if subject not in SUBJECTS:
            print(f"未知科目: {subject}")
            continue
        
        titles = SUBJECTS[subject]
        out_file = CORPUS_DIR / f"{subject}.txt"
        existing = get_existing_titles(out_file)
        
        new_titles = [t for t in titles if t not in existing]
        skip_count = len(existing)
        print(f"\n{'='*60}")
        print(f"[{subject}] 共 {len(titles)} 词条，已有 {skip_count}，新增 {len(new_titles)}")
        print(f"{'='*60}")
        
        # 读取已有内容（追加模式）
        if out_file.exists():
            all_content = out_file.read_text(encoding='utf-8')
        else:
            all_content = f"=== {subject} ===\n"
        
        for i, title in enumerate(new_titles):
            print(f"  [{i+1}/{len(new_titles)}] {title}...", end=' ', flush=True)
            text = fetch_extract(title)
            if text:
                all_content += f"\n\n## {title}\n{text}\n"
                print(f"{len(text)} 字符")
                total_new += 1
            else:
                print("无")
            # 增量保存（每抓 5 条保存一次，防中断丢失）
            if (i + 1) % 5 == 0:
                out_file.write_text(all_content, encoding='utf-8')
            time.sleep(3)  # 长延时防限流
        
        out_file.write_text(all_content, encoding='utf-8')
        print(f"[{subject}] 完成，已保存到 {out_file.name}")
    
    print(f"\n{'='*60}")
    print(f"全部完成：新增 {total_new} 词条，跳过 {total_skip} 已有")
    print(f"语料目录: {CORPUS_DIR}")


if __name__ == '__main__':
    main()
