# -*- coding: utf-8 -*-
"""AI 大模型模板匹配 · 批量执行器。

用法：
  python run_tests.py group1                  # 按 manifest 分组的定向用例（每组一个批次）
  python run_tests.py group1 --upload-only    # 只上传，不发起智能清洗（用于 ALL-01 快照）
  python run_tests.py group1 --limit 10
  python run_tests.py list                    # 列出分组
"""
import argparse
import json
import os
import sys
import time
from collections import defaultdict
from datetime import datetime

import psycopg2
import requests

BASE = os.path.dirname(os.path.abspath(__file__))
API = "http://127.0.0.1:8980/js/a"
DB = dict(host="192.168.0.106", port=5432, dbname="zsjz-ai", user="postgres", password="postgres")
CASE_ID = 20
CREDS = {"username": "admin", "password": "admin123"}

sys.stdout.reconfigure(encoding="utf-8")


def log(msg):
    print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)


class Client:
    def __init__(self):
        self.s = requests.Session()
        self.login()

    def login(self):
        r = self.s.post(API + "/auth/login", json=CREDS, timeout=30)
        r.raise_for_status()
        self.s.headers["x-token"] = r.json()["data"]["token"]

    def open_case(self):
        last = None
        for _ in range(30):
            r = self.s.post(API + "/case/open", json={"id": CASE_ID}, timeout=60)
            last = r
            if r.json().get("code") == 200:
                return
            time.sleep(2)
        raise RuntimeError(f"open case failed: {last.text}")

    def request(self, method, path, **kw):
        for attempt in range(2):
            r = self.s.request(method, API + path, timeout=kw.pop("timeout", 120), **kw)
            try:
                code = r.json().get("code")
            except Exception:
                code = None
            if code == 401 and attempt == 0:
                self.login()
                self.open_case()
                continue
            return r
        return r

    def upload(self, path, batch_id=None, reclean=0):
        params = {}
        if batch_id:
            params["batchId"] = batch_id
        if reclean:
            params["reClean"] = reclean
        with open(path, "rb") as f:
            r = self.request("POST", "/dm/preFileUpload", params=params,
                             files={"file": (os.path.basename(path), f)}, timeout=120)
        return r.json()

    def smart_clean(self, roots):
        r = self.request("POST", "/aiClean/smartClean",
                         json={"fileInfos": roots, "reClean": 0}, timeout=120)
        return r.json()

    def status(self, batch_id):
        r = self.request("GET", "/aiClean/smartClean/status", params={"batchId": batch_id}, timeout=60)
        return (r.json() or {}).get("data")

    def table_page(self, table_name, page=1, limit=100):
        r = self.request("POST", "/gt/treeTablePage",
                         json={"tableName": table_name, "page": page, "limit": limit}, timeout=60)
        return (r.json() or {}).get("data") or {}


def db_jobs(sheet_ids):
    if not sheet_ids:
        return {}
    conn = psycopg2.connect(**DB)
    cur = conn.cursor()
    cur.execute("""select id, sheet_id, status, exit_code, matched_by, match_tier,
                          ref_template_id, ai_template_id, llm_calls, confidence,
                          token_in, token_out, cost_ms, row_total, row_loaded,
                          row_rejected, row_warned, row_unprocessed, fail_reason,
                          draft, detail, model_name, update_time
                   from ai_clean_job where sheet_id = any(%s)
                   order by id""", ([int(x) for x in sheet_ids],))
    rows = cur.fetchall()
    conn.close()
    keys = ["job_id", "sheet_id", "status", "exit_code", "matched_by", "match_tier",
            "ref_template_id", "ai_template_id", "llm_calls", "confidence",
            "token_in", "token_out", "cost_ms", "row_total", "row_loaded",
            "row_rejected", "row_warned", "row_unprocessed", "fail_reason",
            "draft", "detail", "model_name", "update_time"]
    out = {}
    for r in rows:
        d = dict(zip(keys, r))
        for k in ("confidence", "update_time"):
            if d.get(k) is not None:
                d[k] = str(d[k])
        out[int(d["sheet_id"])] = d
    return out


def db_batch(batch_id):
    conn = psycopg2.connect(**DB)
    cur = conn.cursor()
    cur.execute("""select id, stage, file_count, sheet_total, matched_count, ai_judged_count,
                          existing_hit_count, ai_load_count, need_review_count, failed_count,
                          loaded_rows, rejected_rows, unprocessed_rows, self_built_count,
                          last_error, update_time
                   from ai_clean_batch where batch_id = %s order by id desc limit 1""", (int(batch_id),))
    r = cur.fetchone()
    conn.close()
    if not r:
        return None
    keys = ["id", "stage", "file_count", "sheet_total", "matched_count", "ai_judged_count",
            "existing_hit_count", "ai_load_count", "need_review_count", "failed_count",
            "loaded_rows", "rejected_rows", "unprocessed_rows", "self_built_count",
            "last_error", "update_time"]
    d = dict(zip(keys, r))
    d["update_time"] = str(d["update_time"])
    return d


def db_ai_template(tpl_id):
    if not tpl_id:
        return None
    conn = psycopg2.connect(**DB)
    cur = conn.cursor()
    cur.execute("select id, table_name_cn, headers, status, confidence, match_tier from ai_clean_template where id=%s", (int(tpl_id),))
    t = cur.fetchone()
    fields = []
    if t:
        cur.execute("select column_name, field_name_cn, field_type, required, field_sort from ai_clean_field where ai_template_id=%s order by field_sort", (int(tpl_id),))
        fields = [dict(zip(["column_name", "field_name_cn", "field_type", "required", "field_sort"], r)) for r in cur.fetchall()]
    conn.close()
    if not t:
        return None
    return dict(id=t[0], table_name_cn=t[1], headers=t[2], status=t[3],
                confidence=str(t[4]) if t[4] is not None else None, match_tier=t[5], fields=fields)


def upload_one(c, rec, batch_id):
    path = os.path.join(BASE, rec["path"])
    r = c.upload(path, batch_id)
    data = r.get("data")
    out = dict(test_id=rec["test_id"], file=rec["path"], expected_template_id=rec.get("expected_template_id"),
               template_id=rec.get("template_id"), mutation=rec.get("mutation"), notes=rec.get("notes"),
               upload_code=r.get("code"), upload_message=r.get("message"))
    if not data:
        out["upload"] = "DATA_NULL"
        out["sheets"] = []
        return out, None
    root = data
    sheet_list = []
    for ch in (root.get("children") or []):
        sheet_list.append(dict(
            sheet_id=ch.get("id"), sheet_name=ch.get("sheetName"),
            file_status=ch.get("fileStatus"), template_id=ch.get("templateId"),
            table_name=ch.get("tableName"), line_no=ch.get("lineNo"),
            data_num=ch.get("dataNum"), main_id=ch.get("mainId"),
        ))
    out["upload"] = "OK"
    out["batch_id"] = root.get("batchId")
    out["file_id"] = root.get("id")
    out["sheets"] = sheet_list
    return out, root


def run_group(group, upload_only=False, limit=None, batch_size=12, timeout_s=1200, ids=None):
    manifest = json.load(open(os.path.join(BASE, "manifest.json"), encoding="utf-8"))
    files = [x for x in manifest["files"] if x["category"] == group]
    if ids:
        wanted = {x.strip() for x in ids.split(",") if x.strip()}
        files = [x for x in files if x["test_id"] in wanted]
    if limit:
        files = files[:limit]
    if not files:
        log(f"group {group} 没有文件")
        return
    c = Client()
    c.open_case()
    os.makedirs(os.path.join(BASE, "results"), exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_path = os.path.join(BASE, "results", f"{group}_{stamp}.json")
    results = []
    log(f"group={group} files={len(files)} upload_only={upload_only}")

    for start in range(0, len(files), batch_size):
        chunk = files[start:start + batch_size]
        batch_id = None
        roots = []
        chunk_recs = []
        log(f"--- chunk {start + 1}-{start + len(chunk)} ---")
        for rec in chunk:
            if rec.get("format") == "txt" or (rec.get("mutation") or {}).get("_corrupt"):
                # 仍然上传，验证失败节点
                pass
            try:
                out, root = upload_one(c, rec, batch_id)
            except Exception as e:
                out, root = {"test_id": rec["test_id"], "file": rec["path"],
                             "upload": f"EXCEPTION:{type(e).__name__}:{e}"}, None
            if root is not None:
                roots.append(root)
                batch_id = batch_id or root.get("batchId")
            chunk_recs.append(out)
            st = out.get("sheets") or []
            kinds = [(s["sheet_name"], s["file_status"], s["template_id"]) for s in st]
            log(f"  {rec['test_id']}: {out['upload']} {kinds}")
            time.sleep(0.2)

        batch_status = None
        if not upload_only and roots:
            sc = c.smart_clean(roots)
            if sc.get("code") != 200 or not sc.get("data"):
                log(f"  smartClean rejected: {sc.get('code')} {sc.get('message')}")
                for out in chunk_recs:
                    out["smart_clean"] = sc
            else:
                bid = sc["data"].get("batchId") if isinstance(sc.get("data"), dict) else None
                out_batch = bid or batch_id
                log(f"  smartClean accepted batch={out_batch}")
                t0 = time.time()
                last = None
                while time.time() - t0 < timeout_s:
                    st = c.status(out_batch)
                    if st and st != last:
                        last = st
                        log(f"    stage={st.get('stage')} judged={st.get('aiJudgedCount')}/{max(0, (st.get('sheetTotal') or 0) - (st.get('matchedCount') or 0))} "
                            f"needReview={st.get('needReviewCount')} failed={st.get('failedCount')} loaded={st.get('loadedRows')}")
                    if st and st.get("stage") in ("DONE", "PARTIAL", "FAILED"):
                        batch_status = st
                        break
                    time.sleep(3)
                if not batch_status:
                    batch_status = {"stage": "TIMEOUT"}
                log(f"  batch final: {batch_status}")
                for out in chunk_recs:
                    out["smart_clean"] = sc
                    out["batch_final"] = batch_status

        # 回填 job / sheet 结果
        sheet_ids = [s["sheet_id"] for out in chunk_recs for s in (out.get("sheets") or []) if s.get("sheet_id")]
        jobs = db_jobs(sheet_ids)
        for out in chunk_recs:
            out["jobs"] = []
            for s in out.get("sheets") or []:
                job = jobs.get(int(s["sheet_id"])) if s.get("sheet_id") else None
                s["job"] = job
                if job:
                    out["jobs"].append(job)
            # AI 模板元数据
            tpl_ids = {j.get("ai_template_id") for j in out["jobs"] if j.get("ai_template_id")}
            out["ai_templates"] = [db_ai_template(t) for t in tpl_ids]
        results.extend(chunk_recs)
        if batch_status:
            results[-len(chunk_recs)]["chunk_batch"] = batch_status
        json.dump(results, open(out_path, "w", encoding="utf-8"), ensure_ascii=False, indent=1, default=str)
        log(f"  saved {out_path} ({len(results)})")

    # 汇总
    def cnt(pred):
        return sum(1 for x in results if pred(x))
    log("====== summary ======")
    log(f"files={len(results)} data_null={cnt(lambda x: x.get('upload') == 'DATA_NULL')} "
        f"sheet_fail={cnt(lambda x: any(s.get('file_status') == 'SHEET_TEMPLATE_MATCHED_FAIL' for s in x.get('sheets') or []))} "
        f"sheet_ok={cnt(lambda x: any(s.get('file_status') == 'SHEET_TEMPLATE_MATCHED_SUCCESS' for s in x.get('sheets') or []))}")
    jobs = [j for x in results for j in x.get("jobs") or []]
    log(f"jobs={len(jobs)} exits={ {k: sum(1 for j in jobs if j.get('exit_code') == k) for k in ['A','B',None]} } "
        f"by={ {k: sum(1 for j in jobs if j.get('matched_by') == k) for k in ['HEADER_CACHE','SCORE','LLM','SELF_HEADER',None]} }")
    log(f"saved -> {out_path}")


def main():
    global CASE_ID
    ap = argparse.ArgumentParser()
    ap.add_argument("group", nargs="?")
    ap.add_argument("--upload-only", action="store_true")
    ap.add_argument("--limit", type=int)
    ap.add_argument("--batch-size", type=int, default=12)
    ap.add_argument("--timeout", type=int, default=1200)
    ap.add_argument("--case", type=int, default=CASE_ID)
    ap.add_argument("--ids", type=str)
    ap.add_argument("--list", action="store_true")
    args = ap.parse_args()
    CASE_ID = args.case
    manifest = json.load(open(os.path.join(BASE, "manifest.json"), encoding="utf-8"))
    groups = defaultdict(int)
    for x in manifest["files"]:
        groups[x["category"]] += 1
    if args.list or not args.group:
        for k, v in sorted(groups.items()):
            print(f"{k}\t{v}")
        return
    run_group(args.group, args.upload_only, args.limit, args.batch_size, args.timeout, args.ids)


if __name__ == "__main__":
    main()