"""20 题批量问答测试：调 /api/analyst/chat，记录 status/SQL/rows/answer。

需本地 uvicorn + MySQL + DeepSeek Key（真实 LLM）。
用法：python scripts/dev/run_query_battery.py
"""
from __future__ import annotations

import json
import os
import sys
import urllib.request

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, ROOT)

from app.service.auth_service import issue_dev_token

URL = "http://127.0.0.1:8000/api/analyst/chat"

QUERIES = [
    ("Q01", "各风险等级（C1~C5）客户各有多少？占比多少？"),
    ("Q02", "C3 客户大概占多少？"),
    ("Q03", "高净值 vs 非高净值客户各多少、占比多少？"),
    ("Q04", "客户总体规模（总客户数）怎样？"),
    ("Q05", "近 30 天业务活跃度和再往前 30 天比怎样？按交易笔数"),
    ("Q06", "各服务等级（普通/VIP/钻石）客户分布？"),
    ("Q07", "各城市/地区客户大概怎么分布？汇总"),
    ("Q08", "各风险等级产品（R1~R5）各有多少只？"),
    ("Q09", "钱主要集中在哪类产品？按持仓市值"),
    ("Q10", "最近 30 天申购金额和笔数，和再往前 30 天比怎样？"),
    ("Q11", "近 30 天赎回大概多少？"),
    ("Q12", "哪几类产品申购最猛？按产品类型汇总，不要列单个客户"),
    ("Q13", "还有多少条预警没处理（待审核）？"),
    ("Q14", "预警按状态怎么分布？"),
    ("Q15", "高等级预警（反洗钱/适当性）大概多少？"),
    ("Q16", "大额交易类预警大概多少？"),
    ("Q17", "近一段时间适当性不匹配/被拦大概多少笔？汇总"),
    ("Q18", "理财顾问名下客户规模对比"),
    ("Q19", "近 30 天新开户大概多少？"),
    ("Q20", "AML 风险等级（高/中/低）客户占比？"),
]


def call(question: str, idx: int, token: str) -> dict:
    payload = json.dumps({
        "question": question,
        "session_id": f"battery-{idx}",
        "trace_id": f"trace-{idx}",
    }).encode("utf-8")
    req = urllib.request.Request(
        URL,
        data=payload,
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json; charset=utf-8",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=180) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main() -> None:
    token = issue_dev_token(sub="STAFF-20001", roles=["analyst"])
    out = []
    for idx, (tag, q) in enumerate(QUERIES, 1):
        rec: dict = {"tag": tag, "question": q}
        try:
            r = call(q, idx, token)
        except Exception as exc:
            rec.update({"status": "EXC", "error": str(exc)})
            out.append(rec)
            print(json.dumps(rec, ensure_ascii=False))
            continue
        t = r.get("table") or {}
        rec["status"] = r.get("status")
        rec["answer"] = r.get("answer")
        rec["sql"] = (r.get("sql") or "").replace("\n", " ")
        rec["columns"] = t.get("columns")
        rec["rows"] = (t.get("rows") or [])[:6]
        rec["row_count"] = len(t.get("rows") or [])
        rec["error_code"] = r.get("error_code")
        out.append(rec)
        print(json.dumps(rec, ensure_ascii=False))
        sys.stdout.flush()

    report_path = os.path.join(ROOT, "scripts", "dev", "battery_report.json")
    with open(report_path, "w", encoding="utf-8") as f:
        json.dump(out, f, ensure_ascii=False, indent=2)
    print(f"\nWrote {report_path}")


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    main()