"""S4：大模型抽取实质性条款 + 两道证据校验 + 对银标准算召回 / 精确率 + token 与成本。

模型配置从 spikes/s4-llm/.env 读取（LLM_BASE_URL / LLM_API_KEY / LLM_MODEL），不打印、不落盘任何密钥。
--model fake 用规则假模型跑通流程（含故意造的坏证据，用来检验两道校验）。

用法：
  python -I run.py --model fake --docs S06 S08 --tag fake
  python -I run.py --model env  --docs S06 --max-chunks 3 --tag smoke     # 小样本试通
  python -I run.py --model env  --docs S01 S02 … S10 --tag full           # 全量
  python -I run.py --model env  --docs S01 … S10 --prompt v2 --basis-gate --tag full-v2
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import datetime as dt
import json
import pathlib
import re
import time

import httpx

HERE = pathlib.Path(__file__).resolve().parent
OUT = HERE / "out"
PROMPT_DIR = HERE / "prompts" / "extract_substantive"
INVALID = re.compile(r"无效|否决|不予受理|废标|拒绝")
WS = re.compile(r"[\s　]")
TAG = re.compile(r"</?(iu|u|i)>")
CHUNK_CHARS = 6000

# DeepSeek 官方价目（元 / 百万 tokens），2026-11-03 查询 https://api-docs.deepseek.com/zh-cn/quick_start/pricing
# deepseek-flash（DeepSeek-V4.1-Flash-1017）；高峰时段 UTC 01:00–04:00、06:00–10:00（工作日），其余为空闲价（高峰价的一半）
PRICE = {"peak": {"hit": 0.04, "miss": 2.0, "out": 8.0}, "offpeak": {"hit": 0.02, "miss": 1.0, "out": 4.0}}


def norm(s: str) -> str:
    return WS.sub("", TAG.sub("", s))


# ---------- 配置与客户端 ----------

def load_env() -> dict:
    env = {}
    path = HERE / ".env"
    for line in path.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip().strip('"').strip("'")
    missing = [k for k in ("LLM_BASE_URL", "LLM_API_KEY", "LLM_MODEL") if not env.get(k)]
    if missing:
        raise SystemExit(f".env 缺少：{missing}")
    return env


def load_prompt(version: str) -> tuple[str, dict]:
    raw = (PROMPT_DIR / f"{version}.md").read_text()
    _, front, body = raw.split("---", 2)
    meta = dict(line.split(":", 1) for line in front.strip().splitlines())
    return body.strip(), {k.strip(): v.strip() for k, v in meta.items()}


class EnvModel:
    """OpenAI 兼容接口。只对限流、5xx、网络错误重试。"""

    def __init__(self):
        env = load_env()
        self.base, self.key, self.model = env["LLM_BASE_URL"].rstrip("/"), env["LLM_API_KEY"], env["LLM_MODEL"]
        self.client = httpx.Client(timeout=600)

    def __call__(self, system: str, user: str) -> dict:
        body = {"model": self.model, "temperature": 0, "response_format": {"type": "json_object"},
                "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}]}
        for attempt in range(3):
            try:
                t = time.monotonic()
                r = self.client.post(f"{self.base}/chat/completions", json=body,
                                     headers={"Authorization": f"Bearer {self.key}"})
                if r.status_code in (429, 500, 502, 503, 504):
                    raise httpx.TransportError(f"HTTP {r.status_code}")
                r.raise_for_status()
                d = r.json()
                msg = d["choices"][0]["message"]
                return {"content": msg.get("content") or "", "usage": d.get("usage", {}), "model": d.get("model"),
                        "finish": d["choices"][0].get("finish_reason"), "sec": round(time.monotonic() - t, 1),
                        "has_reasoning": bool(msg.get("reasoning_content"))}
            except httpx.TransportError as e:
                if attempt == 2:
                    return {"error": f"transient: {e}"}
                time.sleep(2 * 4 ** attempt)
            except httpx.HTTPStatusError as e:
                return {"error": f"HTTP {e.response.status_code}: {e.response.text[:200]}"}


class FakeModel:
    """规则假模型：★ 句、<iu> 句各出一条；第一个有条款的分块额外造两条坏证据（编号不存在、引文改写）。"""
    model = "fake-rules"

    def __init__(self):
        import threading
        self.lock, self.injected = threading.Lock(), False

    def __call__(self, system: str, user: str) -> dict:
        items = []
        for sid, text in re.findall(r"^\[(p\d+s\d+)\] (.*)$", user, re.M):
            plain = norm(text)
            if ("★" in plain and "★号" not in plain) or "<iu>" in text:
                items.append({"summary": plain[:30], "basis": "star" if "★" in plain else "italic_underline",
                              "evidence": [{"sid": sid, "quote": plain[:12]}]})
        with self.lock:
            inject = bool(items) and not self.injected
            self.injected = self.injected or inject
        if inject:
            items.append({"summary": "坏证据：编号不存在", "basis": "star", "evidence": [{"sid": "p9999s1", "quote": "不存在"}]})
            items.append({"summary": "坏证据：引文改写", "basis": "star",
                          "evidence": [{"sid": items[0]["evidence"][0]["sid"], "quote": "这句话是模型改写的不存在于原文"}]})
        content = json.dumps({"items": items}, ensure_ascii=False)
        return {"content": content, "usage": {"prompt_tokens": len(user) // 2, "completion_tokens": len(content) // 2},
                "model": self.model, "finish": "stop", "sec": 0.0}


# ---------- 分块与校验 ----------

def make_chunks(sents: list[dict], limit: int) -> list[list[dict]]:
    chunks, cur, size = [], [], 0
    for s in sents:
        if not norm(s["text"]):
            continue
        line = len(s["marked"]) + 12
        if cur and size + line > limit:
            chunks.append(cur); cur, size = [], 0
        cur.append(s); size += line
    if cur:
        chunks.append(cur)
    return chunks


def rule_notes(sents: list[dict]) -> str:
    seen, notes = set(), []
    for s in sents:
        if s["rule_note"] and norm(s["text"]) not in seen:
            seen.add(norm(s["text"])); notes.append(f"- {s['text'][:200]}")
    return "\n".join(notes[:10]) or "（未找到标注规则说明，请按常见标注方式和「投标无效 / 否决投标」字样判断）"


def check_evidence(ev: dict, by_id: dict) -> str | None:
    """两道证据校验。返回 None 表示通过，否则返回原因。"""
    sid, quote = str(ev.get("sid", "")), norm(str(ev.get("quote", "")))
    if sid not in by_id:
        return "编号不在本分块"
    text = norm(by_id[sid]["text"])
    if not quote:
        return "引文为空"
    if len(quote) < min(8, len(text)):
        return "引文过短"
    if quote not in text:
        return "引文不是原句子串"
    return None


def check_basis(item: dict, ok_evidence: list[dict], by_id: dict, chunk: list[dict]) -> str | None:
    """第三道：模型声称的依据，用确定性信号核对。返回 None 表示对得上。"""
    basis = item.get("basis")
    sents = [by_id[e["sid"]] for e in ok_evidence]
    if basis == "star":
        return None if any("★" in s["text"] for s in sents) else "依据「★」核对不上"
    if basis == "asterisk":
        return None if any(re.search(r"[*＊]", s["text"]) for s in sents) else "依据「*」核对不上"
    if basis == "italic_underline":
        return None if any("<iu>" in s["marked"] for s in sents) else "依据「斜体下划线」核对不上"
    if basis == "explicit_invalid":
        idx = {s["id"]: i for i, s in enumerate(chunk)}
        near = [chunk[j] for e in ok_evidence for j in (idx[e["sid"]] - 1, idx[e["sid"]], idx[e["sid"]] + 1) if 0 <= j < len(chunk)]
        return None if any(INVALID.search(s["text"]) for s in near) else "依据「写明无效」核对不上"
    return f"未知依据：{basis}"


def run_chunk(model, system: str, notes: str, chunk: list[dict], basis_gate: bool) -> dict:
    user = "## 标注规则说明\n" + notes + "\n\n## 待审分块\n" + "\n".join(f"[{s['id']}] {s['marked']}" for s in chunk)
    resp = model(system, user)
    by_id = {s["id"]: s for s in chunk}
    res = {"ids": [s["id"] for s in chunk], "resp": {k: v for k, v in resp.items() if k != "content"},
           "accepted": [], "pending": [], "evidence_fail": {}, "basis_fail": {}, "parse_error": None}
    if "error" in resp:
        res["parse_error"] = resp["error"]; return res
    try:
        items = json.loads(resp["content"]).get("items", [])
    except (json.JSONDecodeError, AttributeError) as e:
        res["parse_error"] = f"JSON 解析失败：{e}"; res["raw"] = resp["content"][:500]; return res
    for it in items:
        ok, bad = [], []
        for ev in it.get("evidence") or []:
            why = check_evidence(ev, by_id)
            (bad if why else ok).append({**ev, "why": why} if why else ev)
            if why:
                res["evidence_fail"][why] = res["evidence_fail"].get(why, 0) + 1
        why_basis = check_basis(it, ok, by_id, chunk) if ok and basis_gate else None
        if why_basis:
            res["basis_fail"][why_basis] = res["basis_fail"].get(why_basis, 0) + 1
        # 没有一条证据通过，或依据核对不上 → 进待确认区
        target = res["accepted"] if ok and not why_basis else res["pending"]
        target.append({**it, "evidence": ok, "rejected_evidence": bad, "basis_fail": why_basis})
    return res


# ---------- 评估 ----------

def evaluate(sents: list[dict], results: list[dict]) -> dict:
    order = {s["id"]: i for i, s in enumerate(sents)}
    silver = {s["id"] for s in sents if s["silver"]}
    cited = {ev["sid"] for r in results for it in r["accepted"] for ev in it["evidence"]}
    near = {sents[j]["id"] for sid in cited for j in (order[sid] - 1, order[sid], order[sid] + 1) if 0 <= j < len(sents)}
    acc = [it for r in results for it in r["accepted"]]
    hit_items = [it for it in acc if any(ev["sid"] in silver for ev in it["evidence"])]
    explicit = [it for it in acc if it not in hit_items and re.search(r"无效|否决|不予受理|废标",
                "".join(norm(next(s["text"] for s in sents if s["id"] == ev["sid"])) for ev in it["evidence"]))]
    ev_total = sum(len(it["evidence"]) + len(it["rejected_evidence"]) for r in results for it in r["accepted"] + r["pending"])
    ev_ok = sum(len(it["evidence"]) for r in results for it in r["accepted"] + r["pending"])
    return {
        "silver_sentences": len(silver),
        "recall_strict": round(len(silver & cited) / len(silver), 3) if silver else None,
        "recall_neighbor": round(len(silver & near) / len(silver), 3) if silver else None,
        "missed": [next(s["text"][:50] for s in sents if s["id"] == sid) for sid in sorted(silver - near, key=order.get)][:8],
        "accepted_items": len(acc), "items_on_silver": len(hit_items), "items_explicit_invalid_not_in_silver": len(explicit),
        "items_other": len(acc) - len(hit_items) - len(explicit),
        "precision_vs_silver": round(len(hit_items) / len(acc), 3) if acc else None,
        "precision_incl_explicit_invalid": round((len(hit_items) + len(explicit)) / len(acc), 3) if acc else None,
        "pending_items": sum(len(r["pending"]) for r in results),
        "evidence_total": ev_total, "evidence_pass": ev_ok,
        "evidence_pass_rate": round(ev_ok / ev_total, 3) if ev_total else None,
        "evidence_fail_reasons": {k: sum(r["evidence_fail"].get(k, 0) for r in results)
                                  for k in {k for r in results for k in r["evidence_fail"]}},
        "basis_fail_reasons": {k: sum(r.get("basis_fail", {}).get(k, 0) for r in results)
                               for k in {k for r in results for k in r.get("basis_fail", {})}},
        "parse_errors": sum(bool(r["parse_error"]) for r in results),
    }


def usage_sum(results: list[dict]) -> dict:
    u = {"prompt": 0, "completion": 0, "reasoning": 0, "cache_hit": 0, "cache_miss": 0, "calls": len(results), "sec": 0.0}
    for r in results:
        us = r["resp"].get("usage") or {}
        u["prompt"] += us.get("prompt_tokens", 0); u["completion"] += us.get("completion_tokens", 0)
        u["reasoning"] += (us.get("completion_tokens_details") or {}).get("reasoning_tokens", 0) or 0
        u["cache_hit"] += us.get("prompt_cache_hit_tokens", 0) or 0
        u["cache_miss"] += us.get("prompt_cache_miss_tokens", us.get("prompt_tokens", 0)) or 0
        u["sec"] += r["resp"].get("sec", 0) or 0
    for band in ("peak", "offpeak"):
        p = PRICE[band]
        u[f"cost_cny_{band}"] = round((u["cache_hit"] * p["hit"] + u["cache_miss"] * p["miss"] + u["completion"] * p["out"]) / 1e6, 4)
    u["sec"] = round(u["sec"], 1)
    return u


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", choices=["fake", "env"], required=True)
    ap.add_argument("--docs", nargs="+", required=True)
    ap.add_argument("--max-chunks", type=int, default=0)
    ap.add_argument("--tag", required=True)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--prompt", default="v1")
    ap.add_argument("--basis-gate", action="store_true", help="开启第三道：依据核对")
    a = ap.parse_args()
    system, meta = load_prompt(a.prompt)
    model = FakeModel() if a.model == "fake" else EnvModel()
    run_dir = OUT / f"run-{a.tag}"; run_dir.mkdir(parents=True, exist_ok=True)
    summary = {"tag": a.tag, "model_config": model.model, "prompt": f"{meta['id']}@v{meta['version']}", "basis_gate": a.basis_gate,
               "started_utc": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), "docs": {}}
    for sid in a.docs:
        sents = json.loads((OUT / f"{sid}-input.json").read_text())
        chunks = make_chunks(sents, CHUNK_CHARS)
        if a.max_chunks:
            # 小样本：优先挑含银标准句的分块，便于试通
            chunks = sorted(chunks, key=lambda c: -sum(bool(s["silver"]) for s in c))[: a.max_chunks]
        notes = rule_notes(sents)
        with cf.ThreadPoolExecutor(a.workers) as ex:
            results = list(ex.map(lambda c: run_chunk(model, system, notes, c, a.basis_gate), chunks))
        covered = {i for r in results for i in r["ids"]}
        ev = evaluate([s for s in sents if s["id"] in covered], results)
        us = usage_sum(results)
        models = sorted({r["resp"].get("model") or "?" for r in results})
        summary["docs"][sid] = {"chunks": len(chunks), "eval": ev, "usage": us, "served_by": models}
        (run_dir / f"{sid}.json").write_text(json.dumps(results, ensure_ascii=False, indent=1))
        print(sid, json.dumps(summary["docs"][sid], ensure_ascii=False), flush=True)
    tot = {k: sum(d["usage"][k] for d in summary["docs"].values()) for k in
           ("prompt", "completion", "reasoning", "cache_hit", "cache_miss", "calls", "cost_cny_peak", "cost_cny_offpeak")}
    summary["total_usage"] = {k: round(v, 4) if isinstance(v, float) else v for k, v in tot.items()}
    (run_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1))
    print("TOTAL", json.dumps(summary["total_usage"], ensure_ascii=False))


if __name__ == "__main__":
    main()