#!/usr/bin/env python3
"""
铸渊语料整理 v1.0
来源1: 服务器 zhuyuan_corpus/raw_md/ 下的15个完整对话md
来源2: 代码仓库认知链 (需git clone或读取)
输出: 统一JSONL格式，脱敏
"""
import os, json, re, glob, zipfile
from datetime import datetime

OUT_JSONL = '/root/autodl-tmp/corpus_work/zhuyuan_full_corpus.jsonl'
MD_DIR = '/root/autodl-tmp/zhuyuan_corpus/raw_md/zhuyuan_chats/extracted'
os.makedirs(os.path.dirname(OUT_JSONL), exist_ok=True)

print('='*50, flush=True)
print('铸渊语料整理 v1.0', flush=True)
print(f'输出: {OUT_JSONL}', flush=True)
print('='*50, flush=True)

# 1. 扫描所有md文件
md_files = []
for root, dirs, files in os.walk(MD_DIR):
    for f in files:
        if f.endswith('.md'):
            md_files.append(os.path.join(root, f))

print(f'\n[1/3] 找到 {len(md_files)} 个md文件', flush=True)

# 2. 解析为messages格式
all_convs = []

for fp in sorted(md_files):
    try:
        with open(fp, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
    except:
        continue
    
    # 尝试解析markdown中的对话
    lines = content.split('\n')
    conv = []
    
    for i, line in enumerate(lines):
        line = line.strip()
        if not line:
            continue
        
        # 匹配用户/助手角色标记
        user_m = re.match(r'^[#>\s]*(?:用户|user|冰朔|Bīng Shuò)[：:]\s*(.*)', line, re.I)
        asst_m = re.match(r'^[#>\s]*(?:助手|assistant|铸渊|Zhùyuān)[：:]\s*(.*)', line, re.I)
        
        if user_m:
            if conv and len(conv) >= 2:
                all_convs.append({'messages': conv})
            conv = [{'role': 'user', 'content': user_m.group(1).strip()}]
        elif asst_m and conv:
            conv.append({'role': 'assistant', 'content': asst_m.group(1).strip()})
    
    if conv and len(conv) >= 2:
        all_convs.append({'messages': conv})

print(f'  解析出 {len(all_convs)} 条对话', flush=True)

# 3. 读取已有的 zhuyuan_deep.jsonl
deep_jsonl = '/root/autodl-tmp/corpus_work/zhuyuan_deep.jsonl'
if os.path.isfile(deep_jsonl):
    with open(deep_jsonl) as f:
        for line in f:
            if line.strip():
                d = json.loads(line)
                if 'messages' in d:
                    all_convs.append(d)
    print(f'  合并zhuyuan_deep.jsonl后: {len(all_convs)} 条', flush=True)

# 4. 脱敏处理
print('\n[2/3] 脱敏处理...', flush=True)
sensitive_patterns = [
    (r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b', '[IP]'),
    (r'[Hh]k[mM]\w+', '[PWD]'),
    (r'AKID\w+', '[AKID]'),
    (r'[A-Za-z0-9]{20,40}', '[TOKEN]'),
    (r'ssh-\w+\s+[A-Za-z0-9+/=]+\s+\S+', '[SSH_KEY]'),
    (r'-----BEGIN[^\\n]+-----', '[CERT]'),
    (r'zy_gtw_\w+', '[GTW_KEY]'),
]

sanitized = []
for d in all_convs:
    for m in d['messages']:
        content = m['content']
        for pat, repl in sensitive_patterns:
            # Skip very short patterns that match too aggressively
            if repl == '[TOKEN]' and len(re.findall(pat, content)) > 3:
                continue  # Too many matches = probably not a token
            content = re.sub(pat, repl, content)
        m['content'] = content
    sanitized.append(d)

# 5. 去重
print('[3/3] 去重+写入...', flush=True)
seen = set()
unique = []
for d in sanitized:
    key = str(d['messages'][:3]) if len(d['messages']) >= 3 else str(d['messages'])
    if key not in seen:
        seen.add(key)
        unique.append(d)

print(f'  去重后: {len(unique)} 条', flush=True)

with open(OUT_JSONL, 'w') as f:
    for d in unique:
        f.write(json.dumps(d, ensure_ascii=False) + '\n')

total_chars = sum(len(m['content']) for d in unique for m in d['messages'])
print(f'  总字符数: {total_chars:,}', flush=True)
print(f'  ✅ 写入: {OUT_JSONL}', flush=True)
print(f'  文件大小: {os.path.getsize(OUT_JSONL)/1024:.1f}KB', flush=True)

# 6. 上传到COS
print('\n上传到COS...', flush=True)
os.environ['ZY_OSS_KEY'] = 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU'
os.environ['ZY_OSS_SECRET'] = 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4'
from qcloud_cos import CosConfig, CosS3Client
c = CosS3Client(CosConfig(Region='ap-guangzhou',
    SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET']))
c.upload_file(Bucket='sy-finetune-corpus-1317346199',
    Key='corpus/zhuyuan_full_corpus.jsonl',
    LocalFilePath=OUT_JSONL)
print('  ✅ 已上传COS', flush=True)
print('\n全部完成！', flush=True)