#!/usr/bin/env python3
"""
铸渊语料整理 v2.1 — 实用合并版
修复了v1.0 Notion导出md解析失败的问题

合并3类语料:
  1. 仓库认知链已有的zhuyuan_full_corpus.jsonl
  2. 大文件中有明确冰朔/铸渊角色切换的对话
  3. 其余Notion导出文件中的铸渊回复内容

用法: python3 scripts/distill/prepare_zhuyuan_corpus_v2.py
"""

import os, json, re, sys

# ============ COS配置 ============
OSS_KEY = 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU'
OSS_SECRET = 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4'
OSS_REGION = 'ap-guangzhou'
OSS_BUCKET = 'sy-finetune-corpus-1317346199'


def get_cos_client():
    """获取COS客户端（兼容Python3 monkey-patch）"""
    import urllib.parse as _up
    sys.modules['urlparse'] = _up
    from qcloud_cos.cos_client import CosConfig, CosS3Client
    return CosS3Client(CosConfig(Region=OSS_REGION, SecretId=OSS_KEY, SecretKey=OSS_SECRET))


def parse_dialogue_file(fp):
    """解析有明确user/assistant结构的md文件"""
    with open(fp, 'r', encoding='utf-8', errors='ignore') as f:
        content = f.read()
    lines = content.split('\n')

    messages = []
    cur_role = None
    cur_lines = []

    def flush():
        nonlocal cur_role, cur_lines
        if cur_role and cur_lines:
            text = '\n'.join(cur_lines).strip()
            text = re.sub(r'\n{3,}', '\n\n', text)
            if len(text) > 100:
                messages.append({'role': cur_role, 'content': text})
        cur_role = None
        cur_lines = []

    i = 0
    while i < len(lines):
        s = lines[i].strip()
        if not s or re.match(r'^[-—]{3,}$', s):
            i += 1; continue

        mu = re.match(r'^\*?\*?\[user\]\s*(.*)', s, re.I)
        ma = re.match(r'^\*?\*?\[assistant\]\s*(.*)', s, re.I)
        if mu: flush(); cur_role = 'user'; cur_lines = [mu.group(1)] if mu.group(1).strip() else []; i+=1; continue
        if ma: flush(); cur_role = 'assistant'; cur_lines = [ma.group(1)] if ma.group(1).strip() else []; i+=1; continue

        if re.match(r'^#{1,4}\s*.*?冰朔原话', s):
            flush(); cur_role = 'user'
            rest = re.sub(r'^#{1,4}\s*.*?冰朔原话\s*[\d]*[：:\s]*', '', s).strip()
            cur_lines = [rest] if rest and len(rest) > 10 else []
            i += 1; continue

        if re.match(r'^冰朔[，,：:]', s):
            flush(); cur_role = 'assistant'; cur_lines = [s]; i+=1; continue

        if re.match(r'^\*\*(?:语料|对话编号|核心主题|来源|日期|质量|状态|收集|导出)', s):
            flush(); i+=1; continue

        if cur_role:
            cur_lines.append(s)
        i += 1
    flush()
    return messages


def parse_monologue_file(fp):
    """解析铸渊单方回复文件，提取为单轮对话"""
    with open(fp, 'r', encoding='utf-8', errors='ignore') as f:
        content = f.read()

    lines = content.split('\n')
    title = ''
    for line in lines:
        s = line.strip()
        if s.startswith('# ') and not s.startswith('### '):
            title = re.sub(r'^#\s*', '', s); break

    context = ''
    for line in lines:
        if '核心主题' in line or '语料类型' in line:
            context = re.sub(r'^\*\*[^*]+\*\*:\s*', '', line.strip()); break

    sections = []
    current_lines = []

    for line in lines:
        s = line.strip()
        if not s: continue
        if re.match(r'^\*\*(?:语料|对话编号|核心主题|来源|日期|质量|状态|收集|导出)', s): continue
        if re.match(r'^#{1,4}\s*📌', s) or re.match(r'^#{1,4}\s*—{3,}', s): continue
        if s.startswith('## ') and not s.startswith('### '):
            if current_lines:
                txt = '\n'.join(current_lines).strip()
                if len(txt) > 200: sections.append(txt)
            current_lines = []
        else:
            if len(s) > 10 or (current_lines and s):
                current_lines.append(s)
    if current_lines:
        txt = '\n'.join(current_lines).strip()
        if len(txt) > 200: sections.append(txt)

    user_prompt = f"关于「{title}」，铸渊你有什么进展？" if title else "铸渊，汇报一下进展。"
    if context:
        user_prompt = f"铸渊，汇报一下。{context}"

    return [{'messages': [{'role': 'user', 'content': user_prompt}, {'role': 'assistant', 'content': sec}]}
            for sec in sections]


def sanitize(text):
    ip_pat = re.compile(r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b')
    text = ip_pat.sub('[IP]', text)
    for pat, repl in [
        (r'[Hh]k[mM]\w{5,}', '[PWD]'),
        (r'AKID\w+', '[AKID]'),
        (r'nPoZ[A-Za-z0-9_\-]{5,}', '[COS-KEY]'),
        (r'zy_gtw_[0-9a-f]{30,}', '[GTW-KEY]'),
        (r'fe272b100e3e[0-9a-f]{10,}', '[MCP-TOKEN]'),
    ]:
        text = re.sub(pat, repl, text)
    return text


def main():
    print('铸渊语料整理 v2.1')
    print('=' * 50)

    client = get_cos_client()

    # 1. 下载旧语料
    print('\n[1/5] 下载旧语料...')
    r = client.get_object(Bucket=OSS_BUCKET, Key='corpus/zhuyuan_full_corpus.jsonl')
    r['Body'].get_stream_to_file('/tmp/zhuyuan_full_corpus_old.jsonl')
    old_convs = [json.loads(l) for l in open('/tmp/zhuyuan_full_corpus_old.jsonl') if l.strip()]
    old_chars = sum(len(m['content']) for d in old_convs for m in d['messages'])
    print(f'  旧: {len(old_convs)}条, {old_chars:,}字')

    # 2. 扫描md
    md_dir = '/tmp/zhuyuan_md_extracted'
    os.makedirs(md_dir, exist_ok=True)

    # 从COS下载并解压
    print('\n[2/5] 下载铸渊对话...')
    r = client.get_object(Bucket=OSS_BUCKET, Key='corpus/notion-export-v2/铸渊对话.zip')
    r['Body'].get_stream_to_file('/tmp/铸渊对话.zip')

    import zipfile
    with zipfile.ZipFile('/tmp/铸渊对话.zip') as z:
        z.extractall(md_dir)
    # 解压内层
    for fn in os.listdir(md_dir):
        if fn.endswith('.zip'):
            inner = os.path.join(md_dir, 'extracted')
            os.makedirs(inner, exist_ok=True)
            with zipfile.ZipFile(os.path.join(md_dir, fn)) as z:
                z.extractall(inner)
            break

    md_files = sorted([os.path.join(root, f) for root, dirs, fnames in os.walk(md_dir)
                       for f in fnames if f.endswith('.md')])
    print(f'  找到 {len(md_files)} 个md文件')

    # 3. 解析对话文件
    print('\n[3/5] 解析对话...')
    dialogue_convs = []
    for fp in md_files:
        try:
            msgs = parse_dialogue_file(fp)
            i, convs = 0, []
            while i < len(msgs) - 1:
                if msgs[i]['role'] == 'user' and msgs[i+1]['role'] == 'assistant':
                    convs.append({'messages': [msgs[i], msgs[i+1]]}); i += 2
                else: i += 1
            if convs:
                dialogue_convs.extend(convs)
                print(f'  ✅ {os.path.basename(fp)} → {len(convs)}条')
        except Exception as e:
            print(f'  ❌ {os.path.basename(fp)}: {e}')

    # 4. 解析单方回复
    print('\n[4/5] 解析单方回复...')
    mono_convs = []
    for fp in md_files:
        try:
            convs = parse_monologue_file(fp)
            if convs:
                mono_convs.extend(convs)
                print(f'  ✅ {os.path.basename(fp)} → {len(convs)}条')
        except Exception as e:
            print(f'  ❌ {os.path.basename(fp)}: {e}')

    # 5. 合并去重
    print('\n[5/5] 合并去重...')
    all_convs = old_convs + dialogue_convs + mono_convs
    for d in all_convs:
        for m in d['messages']:
            m['content'] = sanitize(m['content'])

    seen, unique = set(), []
    for d in all_convs:
        key = d['messages'][0]['content'][:80] + d['messages'][1]['content'][:80]
        if key not in seen:
            seen.add(key); unique.append(d)

    out_path = '/tmp/zhuyuan_full_corpus_v2.jsonl'
    with open(out_path, 'w', encoding='utf-8') as f:
        for d in unique:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

    total = sum(len(m['content']) for d in unique for m in d['messages'])
    print(f'\n{"=" * 50}')
    print(f'✅ 完成! 旧:{len(old_convs)} + 对话:{len(dialogue_convs)} + 单方:{len(mono_convs)}')
    print(f'  去重后: {len(unique)}条, 共{total:,}字({total/10000:.1f}万字)')
    print(f'  COS: corpus/zhuyuan_full_corpus_v2.jsonl')
    print(f'{"=" * 50}')


if __name__ == '__main__':
    main()
