#!/usr/bin/env python3
"""
高温度多样化生成 — 单线程，与常规跑批共用 Astron 5并发中的1个
对已完成批1-6的问题，用高温度重新生成，获得多样回答
"""
import sys, os, json, random, time, requests
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)

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"
MODEL_DIR = os.path.dirname(os.path.abspath(__file__))

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

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

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


def call_astron(q, temperature):
    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": 8192,
            }
            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:
                    return content
            else:
                time.sleep(1)
        except Exception:
            time.sleep(2)
    return None


def main():
    # 从已完成批载入问题
    all_questions = []
    for i in range(1, 7):
        fp = os.path.join(MODEL_DIR, f'mega_interim_{i}.json')
        if os.path.exists(fp):
            with open(fp) as f:
                items = json.load(f)
            for item in items:
                all_questions.append(item['question'])
            print(f"批{i}: {len(items)}条已载入")
    
    # 去重
    all_questions = list(dict.fromkeys(all_questions))
    print(f"共 {len(all_questions)} 个去重问题")
    
    # 温度配置
    temp_configs = [0.9, 1.3, 1.5]
    all_new = []
    
    for round_idx, temp in enumerate(temp_configs):
        print(f"\n{'='*50}")
        print(f"轮{round_idx+1}/3: 温度={temp}")
        print(f"{'='*50}")
        
        # 单线程顺序跑
        results = []
        total = len(all_questions)
        t0 = time.time()
        
        for idx, q in enumerate(all_questions):
            ans = call_astron(q, temp)
            if ans:
                results.append({"question": q, "answer": ans})
            
            if (idx + 1) % 20 == 0 or idx + 1 == total:
                elapsed = max(time.time() - t0, 1)
                rate = len(results) / elapsed * 60
                eta = (total - idx - 1) / max(idx + 1, 1) * elapsed / 60
                print(f"  多样: {idx+1}/{total} (已获{len(results)}条, {rate:.0f}条/分, ETA:{eta:.0f}分)", end="\r")
        
        elapsed = time.time() - t0
        rate = len(results) / max(elapsed, 1) * 60
        print(f"\n  轮{round_idx+1}完成: {len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)")
        
        for r in results:
            r['temperature'] = temp
            r['round'] = round_idx + 1
        all_new.extend(results)
        
        # ✅ 每轮跑完立即去重并追加到训练集
        seen = set()
        unique = []
        for r in all_new:
            key = (r['question'][:100], r['answer'][:100])
            if key not in seen:
                seen.add(key)
                unique.append(r)
        
        for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']:
            tp = os.path.join(MODEL_DIR, tgt)
            with open(tp, encoding='utf-8') as f:
                cur = json.load(f)
            ex = set((it.get("question","")[:100], it.get("answer","")[:100]) for it in cur)
            to_add = []
            for r in unique:
                key = (r['question'][:100], r['answer'][:100])
                if key not in ex:
                    ex.add(key)
                    to_add.append({"question": r['question'], "answer": r['answer']})
            cur.extend(to_add)
            with open(tp, 'w', encoding='utf-8') as f:
                json.dump(cur, f, ensure_ascii=False, indent=2)
            if 'train' in tgt:
                print(f"✅ 轮{round_idx+1} 训练集 +{len(to_add)}条 → {len(cur)}条")


if __name__ == '__main__':
    main()