#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
单批生成 — 从mega_questions.json取指定batch生成并追加
用法: python3 generate_batch.py <batch_num> <batch_size>
"""
import sys, os, json, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)

from api_correction import batch_generate

model_dir = os.path.dirname(os.path.abspath(__file__))
batch_num = int(sys.argv[1])
batch_size = int(sys.argv[2])

# 加载问题列表
with open(os.path.join(model_dir, 'mega_questions.json'), 'r', encoding='utf-8') as f:
    mega = json.load(f)

# 计算本批范围
start = (batch_num - 1) * batch_size
end = min(start + batch_size, len(mega))
chunk = [item['question'] for item in mega[start:end]]

if not chunk:
    print(f"批{batch_num}: 无数据")
    sys.exit(0)

print(f"批 {batch_num}: {len(chunk)}题 (序号{start}-{end})")
t0 = time.time()

results = batch_generate(chunk, temperature=0.85)
dt = time.time() - t0

# 保存中间文件
interim_path = os.path.join(model_dir, f'mega_interim_{batch_num}.json')
with open(interim_path, 'w', encoding='utf-8') as f:
    json.dump(results, f, ensure_ascii=False, indent=2)

# 追加到训练集
for tgt in ['xiaoli_train_v2.json', 'xiaoli_merged_data_v2.json']:
    tgt_path = os.path.join(model_dir, tgt)
    if not os.path.exists(tgt_path):
        continue
    with open(tgt_path, 'r', encoding='utf-8') as f:
        cur = json.load(f)
    existing = set((item.get("question",""), item.get("answer","")) for item in cur)
    truly_new = [item for item in results if (item["question"], item["answer"]) not in existing]
    if truly_new:
        cur.extend(truly_new)
        with open(tgt_path, 'w', encoding='utf-8') as f:
            json.dump(cur, f, ensure_ascii=False, indent=2)
        if tgt == 'xiaoli_train_v2.json':
            print(f"  ✅ 追加训练集: +{len(truly_new)}条, 现有{len(cur)}条")

print(f"批{batch_num}完成: +{len(results)}条 ({dt:.0f}s)")