#!/usr/bin/env python3
"""llm skill：New-API 统一调用入口（默认 MiniMax-M3）。

管线内所有内容生产型 LLM 调用（清洗探测/拆书抽取）必须经此入口：
- trust_env=False（本机代理环境变量会劫持内网直连，教训固化）；
- 超时 + 指数退避重试；<think> 剥离；JSON 三级容错提取（json-repair 兜底）；
- 每次调用向 stderr 打印 token 用量与耗时（成本审计），stdout 只出内容。
"""
import json
import pathlib
import re
import sys
import time

import click
import requests

BASE = "http://100.64.0.8:3000"
# New-API 普通令牌（仓库政策允许明文；严禁换管理令牌打 /v1）
TOKEN = "sk-DyVqO3lDmEvQZ3PqGpbNaaaHZHhbh0xaHRIiynhYSmVlLHl2"
DEFAULT_MODEL = "MiniMax-M3"


def chat(prompt, model=DEFAULT_MODEL, max_tokens=32000, temperature=0.2,
         retries=2, timeout=900):
    """单轮对话，返回 (content, usage)。网络错/5xx/429 指数退避重试。

    content 已剥离 <think>…</think>（推理模型可能把思考混进正文）。
    """
    s = requests.Session()
    s.trust_env = False  # 本机代理 env 会劫持内网直连
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temperature,
    }
    last_err = None
    for attempt in range(retries + 1):
        try:
            t0 = time.time()
            r = s.post(f"{BASE}/v1/chat/completions",
                       headers={"Authorization": f"Bearer {TOKEN}"},
                       json=payload, timeout=timeout)
            # 429/5xx 属于可重试的服务端瞬时问题
            if r.status_code in (429,) or r.status_code >= 500:
                last_err = f"HTTP {r.status_code}: {r.text[:200]}"
                raise requests.RequestException(last_err)
            r.raise_for_status()
            data = r.json()
            content = data["choices"][0]["message"]["content"] or ""
            content = re.sub(r"<think>.*?</think>", "", content, flags=re.S).strip()
            usage = data.get("usage", {})
            print(f"[llm] {model} in={usage.get('prompt_tokens', '?')} "
                  f"out={usage.get('completion_tokens', '?')} "
                  f"耗时{time.time() - t0:.0f}s finish={data['choices'][0].get('finish_reason')}",
                  file=sys.stderr)
            return content, usage
        except (requests.RequestException, KeyError, json.JSONDecodeError) as e:
            last_err = str(e)
            if attempt < retries:
                wait = 8 * (2 ** attempt)
                print(f"[llm] 第{attempt + 1}次失败({last_err[:120]})，{wait}s 后重试",
                      file=sys.stderr)
                time.sleep(wait)
    raise RuntimeError(f"LLM 调用重试耗尽: {last_err}")


def extract_json(text):
    """JSON 三级容错提取：直接解析 → 首尾括号截取 → json-repair 兜底。

    opus 试拆实测过两类 JSON 病（中文引号、缺逗号）——任何模型都可能犯，统一在此兜住。
    """
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    # 剥 markdown 代码围栏后按最外层大括号/中括号截取
    t = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.M)
    for a, b in (("{", "}"), ("[", "]")):
        i, j = t.find(a), t.rfind(b)
        if i != -1 and j > i:
            frag = t[i:j + 1]
            try:
                return json.loads(frag)
            except json.JSONDecodeError:
                import json_repair
                return json_repair.loads(frag)
    import json_repair
    return json_repair.loads(t)


@click.group()
def cli():
    pass


@cli.command("chat")
@click.option("--prompt-file", type=click.Path(exists=True), required=True)
@click.option("--model", default=DEFAULT_MODEL, show_default=True)
@click.option("--max-tokens", type=int, default=32000, show_default=True)
@click.option("--temperature", type=float, default=0.2, show_default=True)
@click.option("--out", type=click.Path(), help="内容写文件（不给则打 stdout）")
@click.option("--extract-json", "extract_", is_flag=True, help="容错提取 JSON 后输出")
def chat_cmd(prompt_file, model, max_tokens, temperature, out, extract_):
    prompt = pathlib.Path(prompt_file).read_text()
    content, _ = chat(prompt, model=model, max_tokens=max_tokens, temperature=temperature)
    if extract_:
        content = json.dumps(extract_json(content), ensure_ascii=False, indent=1)
    if out:
        pathlib.Path(out).write_text(content)
        click.echo(f"已写 {out}（{len(content):,} 字符）", err=True)
    else:
        click.echo(content)


if __name__ == "__main__":
    try:
        cli()
    except RuntimeError as e:
        click.echo(f"[llm错误] {e}", err=True)
        sys.exit(1)