#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多样化回答生成 — 对每个问题用 Astron 高温生成 50 个不同回答
5并发，若重复则重生成，增量保存防丢
"""
import json, requests, time, os, sys, random, hashlib
from concurrent.futures import ThreadPoolExecutor, as_completed
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)

# ====== Astron 配置 ======
ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions"
ASTRON_KEY = os.environ.get("ASTRON_KEY", "")
ASTRON_MODEL = "astron-code-latest"

SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。

你的人设：
- 8岁女孩子，有点害羞但很可爱
- 说话简单自然，符合小朋友的语气
- 回答问题要简短，1-2句话就够了
- 不要用"喵"结尾，不要过度卖萌

回答要求：
1. 语气自然，像真实的小朋友说话
2. 不懂就说"不知道"
3. 回答简洁但要说清楚，不要讲一半就停"""


def call_astron(q, temperature):
    """调用 Astron，返回回答或 None"""
    for attempt in range(3):
        try:
            headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
            payload = {
                "model": ASTRON_MODEL,
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": q}
                ],
                "temperature": temperature,
                "max_tokens": 2048,
            }
            resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60)
            if resp.status_code == 200:
                content = resp.json()['choices'][0]['message']['content'].strip()
                if content and len(content) > 2:
                    return content
            else:
                time.sleep(1)
        except Exception:
            time.sleep(2)
    return None


def answer_fingerprint(text):
    """回答指纹 — 去空格标点后取前 30 字"""
    clean = ''.join(text.split())[:30]
    return hashlib.md5(clean.encode()).hexdigest()


def is_duplicate(new, existing):
    """检查是否与已有回答重复（精确+模糊）"""
    fp = answer_fingerprint(new)
    if fp in existing:
        return True
    # 模糊：如果新回答是某个已有回答的子串或反之
    short_new = ''.join(new.split())[:20]
    for old_fp, old_text in existing.items():
        short_old = ''.join(old_text.split())[:20]
        if short_new in short_old or short_old in short_new:
            return True if len(short_new) > 8 else False
    return False


def generate_answers_for_question(q, target=50, max_attempts=100, temperature=0.95):
    """为单个问题生成 target 个不同回答"""
    answers = {}  # fp -> full_text
    attempt = 0
    t_start = time.time()
    
    while len(answers) < target and attempt < max_attempts:
        # 温度在 0.85~1.0 之间随机波动，避免卡在同一模式
        temp = temperature + random.uniform(-0.1, 0.05)
        temp = max(0.85, min(1.0, temp))
        
        result = call_astron(f"问题：{q}\n\n请以小狸的身份回答：", temp)
        if result:
            if not is_duplicate(result, answers):
                fp = answer_fingerprint(result)
                answers[fp] = result
        
        attempt += 1
        
        if attempt % 20 == 0:
            elapsed = time.time() - t_start
            rate = len(answers) / max(elapsed, 1)
            if len(answers) < target and rate > 0:
                eta = (target - len(answers)) / rate
            else:
                eta = 0
            print(f"    [{len(answers)}/{target}] attempt={attempt}, {rate:.1f}条/秒, 预计剩余{eta:.0f}秒")
            sys.stdout.flush()
    
    return list(answers.values())


def main():
    model_dir = '/workspace/xiaoli-llm-chat'
    output_dir = os.path.join(model_dir, 'diverse_answers')
    os.makedirs(output_dir, exist_ok=True)
    
    # ====== 参数（可调） ======
    QUESTIONS_PER_BATCH = 100   # 每批处理多少问题
    ANSWERS_PER_QUESTION = 50   # 每个问题生成多少个不同回答
    MAX_ATTEMPTS = 150          # 最多尝试多少次（超过就放弃）
    MAX_WORKERS = 5             # Astron 5并发
    TEMPERATURE = 0.95          # 高温（0.85~1.0 浮动）
    
    # ====== 加载问题列表 ======
    train_path = os.path.join(model_dir, 'xiaoli_train_v2.json')
    with open(train_path, encoding='utf-8') as f:
        data = json.load(f)
    
    # 取所有不重复的问题
    all_questions = list(set(item.get('question', '') for item in data))
    random.shuffle(all_questions)
    print(f"训练集总条数: {len(data)}")
    print(f"不重复问题数: {len(all_questions)}")
    
    # 统计问题出现频次，优先处理高频问题
    q_freq = {}
    for item in data:
        q = item.get('question', '')
        q_freq[q] = q_freq.get(q, 0) + 1
    
    # 按频次排序（高频靠前，低频靠后，随机打乱同频）
    freq_buckets = {}
    for q in all_questions:
        f = q_freq.get(q, 1)
        freq_buckets.setdefault(f, []).append(q)
    
    sorted_questions = []
    for freq in sorted(freq_buckets.keys(), reverse=True):
        bucket = freq_buckets[freq]
        random.shuffle(bucket)
        sorted_questions.extend(bucket)
    
    print(f"按频次排序完成 (高频优先处理)")
    
    # ====== 恢复已有进度 ======
    progress_file = os.path.join(output_dir, 'diversity_progress.json')
    completed = {}  # question -> list of answers
    if os.path.exists(progress_file):
        with open(progress_file, encoding='utf-8') as f:
            completed = json.load(f)
        # 过滤掉已处理的问题
        processed_qs = set(completed.keys())
        sorted_questions = [q for q in sorted_questions if q not in processed_qs]
        print(f"已有进度: {len(completed)} 个问题已完成")
        print(f"剩余未处理: {len(sorted_questions)} 个问题")
    
    # 限制要处理的问题数
    max_q = len(sorted_questions)  # 全部问题
    
    total_questions = len(sorted_questions)
    total_batches = (total_questions + QUESTIONS_PER_BATCH - 1) // QUESTIONS_PER_BATCH
    
    print(f"\n目标: {total_questions} 个问题 × {ANSWERS_PER_QUESTION} 个回答 = {total_questions * ANSWERS_PER_QUESTION} 条")
    print(f"批次: {total_batches} 批 (每批 ≤{QUESTIONS_PER_BATCH} 问题)")
    print(f"并发: {MAX_WORKERS}")
    if total_questions > 1000:
        minutes = total_questions * ANSWERS_PER_QUESTION / (MAX_WORKERS * 8) / 60
        print(f"预估耗时: ~{minutes:.0f} 分钟 (~{minutes/60:.1f} 小时)")
    print(f"{'='*70}\n")
    
    t_start = time.time()
    total_new_pairs = 0
    
    for batch_idx in range(total_batches):
        batch_qs = sorted_questions[batch_idx * QUESTIONS_PER_BATCH : (batch_idx + 1) * QUESTIONS_PER_BATCH]
        batch_results = {}
        
        print(f"\n批次 {batch_idx + 1}/{total_batches} ({len(batch_qs)} 个问题)")
        
        # 并发生成
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = {
                executor.submit(generate_answers_for_question, q, ANSWERS_PER_QUESTION, MAX_ATTEMPTS, TEMPERATURE): q
                for q in batch_qs
            }
            
            for future in as_completed(futures):
                q = futures[future]
                try:
                    answers = future.result()
                    if answers:
                        batch_results[q] = answers
                        print(f"  ✅ 『{q[:30]}』 → {len(answers)} 个不同回答")
                    else:
                        print(f"  ❌ 『{q[:30]}』 → 生成失败")
                except Exception as e:
                    print(f"  ❌ 『{q[:30]}』 → 错误: {e}")
        
        # 追加到总进度
        completed.update(batch_results)
        batch_new_pairs = sum(len(v) for v in batch_results.values())
        total_new_pairs += batch_new_pairs
        
        # 保存进度
        with open(progress_file, 'w', encoding='utf-8') as f:
            json.dump(completed, f, ensure_ascii=False, indent=2)
        
        elapsed = time.time() - t_start
        rate = total_new_pairs / max(elapsed, 1) * 60
        
        # 统计
        print(f"\n批次 {batch_idx + 1} 完成 | "
              f"本批新增 {batch_new_pairs} 条 | "
              f"累计 {total_new_pairs:,} 条 ({len(completed)} 问题) | "
              f"{rate:.0f}条/分 | "
              f"已用 {elapsed/60:.1f}分")
    
    # ====== 完成 ======
    elapsed_total = time.time() - t_start
    print(f"\n{'='*70}")
    print(f"全部生成完成!")
    print(f"处理问题: {len(completed)} 个")
    print(f"总回答数: {total_new_pairs:,} 条")
    print(f"用时: {elapsed_total/60:.1f} 分钟")
    
    # ====== 保存为训练格式 ======
    train_format = []
    for q, answers in completed.items():
        for a in answers:
            train_format.append({'question': q, 'answer': a})
    
    output_file = os.path.join(output_dir, 'diverse_train_data_v1.json')
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(train_format, f, ensure_ascii=False, indent=2)
    
    # 也保存一份到工作目录
    work_file = '/data/user/work/diverse_train_data_v1.json'
    with open(work_file, 'w', encoding='utf-8') as f:
        json.dump(train_format, f, ensure_ascii=False, indent=2)
    
    print(f"\n结果保存:")
    print(f"  📄 {output_file}")
    print(f"  📄 {work_file}")
    print(f"  📊 平均每问题 {total_new_pairs/max(len(completed),1):.1f} 个回答")
    print(f"\n如需合并到训练集：")
    print(f"  python3 -c \"import json; d1=json.load(open('xiaoli_train_v2.json')); d2=json.load(open('{output_file}')); d1.extend(d2); json.dump(d1, open('xiaoli_train_v2.json','w'), ensure_ascii=False, indent=2)\"")
    print(f"{'='*70}")

if __name__ == '__main__':
    main()