#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
API生成模块 — Astron 5并发，稳定不超限
"""
import json, requests, time, os, threading
from concurrent.futures import ThreadPoolExecutor, as_completed

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=0.85):
    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 batch_generate(questions, temperature=0.85, max_workers=4):
    """Astron 4并发 — 明确限制，稳定可靠"""
    total = len(questions)
    results = []
    
    def call(q):
        ans = call_astron(q, temperature)
        if ans:
            return q, ans
        return None
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(call, item): i for i, item in enumerate(questions)}
        done = 0
        t0 = time.time()
        for f in as_completed(futures):
            done += 1
            r = f.result()
            if r:
                results.append({"question": r[0], "answer": r[1]})
            if done % 20 == 0 or done == total:
                elapsed = max(time.time() - t0, 1)
                rate = len(results) / elapsed * 60
                print(f"    Astron: {done}/{total} (已获{len(results)}条, {rate:.0f}条/分)", end="\r")
    
    elapsed = time.time() - t0
    rate = len(results) / max(elapsed, 1) * 60
    print(f"    Astron: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)")
    return results


# 保持旧函数名兼容
batch_generate_astron = batch_generate
batch_generate_mimo = batch_generate
batch_generate_direct = batch_generate
batch_generate_dual = batch_generate


def batch_convert(items, *args, **kwargs):
    """兼容占位：训练脚本导入用，已弃用"""
    print("  [API] batch_convert已弃用，使用batch_generate")
    return []


def batch_generate_multi_temp(questions, *args, **kwargs):
    """兼容占位：训练脚本导入用，已弃用"""
    print("  [API] batch_generate_multi_temp已弃用，使用batch_generate")
    return []