"""造一份示例数据（读者不需要跑，data/ 里已经有现成的）。

做法和 5.11 的文风 LoRA 一样，叫“反向合成”：拿一位作者的原文，让一个大模型用它默认的写法把每段改写一遍，
得到“AI 默认写法 → 作者原文”的配对。训练时输入是 AI 写法，目标是原文。
这里的作者是鲁迅：《朝花夕拾》和《呐喊》都已进入公有领域，可以随书分发。原文取自中文维基文库（简体）。

用法：
    export AI_BUILDER_TOKEN=sk_...
    python make_demo_data.py
"""
import concurrent.futures as cf
import html
import json
import os
import random
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

from openai import OpenAI
from opencc import OpenCC          # pip install opencc-python-reimplemented；维基文库的原文是繁体，转成简体

PAGES = ['狗·猫·鼠', '阿長與《山海經》', '二十四孝圖', '五猖會', '无常', '從百草園到三味書屋', '父亲的病', '琐记',
         '藤野先生', '范愛農', '狂人日記', '孔乙己', '藥', '明天', '一件小事', '頭髮的故事', '風波', '故鄉',
         '端午節', '白光', '兔和貓', '鴨的喜劇', '社戲']
REWRITE = ('把下面这段文字用现代、通顺、规范的书面语重写一遍。意思和信息保持不变，不要增删内容，'
           '不要解释，只输出重写后的这一段。\n\n{text}')
HERE = Path(__file__).resolve().parent
T2S = OpenCC('t2s')
QUOTES = str.maketrans({'『': '“', '』': '”', '「': '‘', '」': '’'})


def fetch(page):
    """取一篇的 HTML（缓存在 data/raw/，维基文库限流，每次请求之间停几秒）。"""
    cache = HERE / 'data' / 'raw' / f'{page}.html'
    if cache.exists():
        body = cache.read_text(encoding='utf-8')
    else:
        url = ('https://zh.wikisource.org/w/api.php?action=parse&prop=text&format=json&variant=zh-cn&page='
               + urllib.parse.quote(page))
        req = urllib.request.Request(url, headers={'User-Agent': 'textbook-demo-data/1.0'})
        for wait in [5, 30, 60, 120]:
            time.sleep(wait)
            try:
                body = json.load(urllib.request.urlopen(req))['parse']['text']['*']
                break
            except urllib.error.HTTPError as e:
                if e.code != 429:
                    raise
        else:
            raise SystemExit('维基文库一直限流，过一会儿再试')
        cache.parent.mkdir(parents=True, exist_ok=True)
        cache.write_text(body, encoding='utf-8')
    paras = []
    for p in re.findall(r'<p>(.*?)</p>', body, re.S):
        t = html.unescape(re.sub(r'<[^>]+>', '', p)).strip()
        t = re.sub(r'\[\d+\]', '', t).replace('\u200b', '')   # 去掉注释编号和零宽空格
        t = T2S.convert(t).translate(QUOTES)
        if 60 <= len(t) <= 320:
            paras.append(t)
    return paras


def main():
    client = OpenAI(base_url='https://space.ai-builders.com/backend/v1', api_key=os.environ['AI_BUILDER_TOKEN'])
    paras = []
    for page in PAGES:
        got = fetch(page)
        print(f'{page}: {len(got)} 段')
        paras += [(page, t) for t in got]
    random.Random(0).shuffle(paras)

    def rewrite(item):
        page, text = item
        r = client.chat.completions.create(model='deepseek', temperature=0.7,
                                           messages=[{'role': 'user', 'content': REWRITE.format(text=text)}])
        return dict(source=page, ai=r.choices[0].message.content.strip(), original=text)

    with cf.ThreadPoolExecutor(8) as ex:
        rows = list(ex.map(rewrite, paras))
    (HERE / 'data').mkdir(exist_ok=True)
    splits = {'test': rows[:20], 'valid': rows[20:40], 'train': rows[40:]}
    for name, rs in splits.items():
        with open(HERE / 'data' / f'pairs_{name}.jsonl', 'w', encoding='utf-8') as f:
            for r in rs:
                f.write(json.dumps(r, ensure_ascii=False) + '\n')
        print(f'{name}: {len(rs)} 对')


if __name__ == '__main__':
    main()