#!/usr/bin/env python3
"""用 jev-1.13 教师模型清洗本地 archive 数据（完整 Jev 请求体 + 批次并发）。

- 输入：data/processed/archive/messages_labeled.jsonl（v1/systemone 判定）
- 模板：.work/teacher/jev_request_template.json（完整 is_ad instructions + 4 questions）
- 每条消息构造 state：chat/message/sender（能拿到的 user_name）/recent_context（同群前 2 条）
- 一次请求带完整 4 questions；客户端 ThreadPool 并发
- 输出：data/cleaned/teacher_<n>.jsonl + report.md/json
"""
import argparse
import json
import random
import re
import threading
import time
import urllib.error
import urllib.request
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from itertools import groupby
from pathlib import Path

LINK_RE = re.compile(r"https?://|t\.me/|www\.", re.I)
USER_RE = re.compile(r"^(.*?)\s*\(@([^)]+)\)")


def load_jsonl(p):
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def load_env(p=".work/teacher/.env"):
    cfg = {}
    for line in Path(p).read_text(encoding="utf-8").splitlines():
        if "=" in line:
            k, v = line.split("=", 1)
            cfg[k.strip()] = v.strip()
    return cfg


def load_judgment_names(root="data/raw/archive/archive"):
    names = {}
    for p in Path(root).glob("*/judgments-*.jsonl"):
        for line in open(p, encoding="utf-8"):
            try:
                r = json.loads(line)
            except Exception:
                continue
            key = (str(r.get("chat_id")), int(r.get("message_id", -1)))
            name = r.get("user_name") or ""
            m = USER_RE.match(name)
            names[key] = (m.group(1).strip(), m.group(2).strip()) if m else (name.strip(), "")
    return names


def make_state(row, names, prev_map):
    text = row.get("text") or row.get("text_raw") or ""
    text = str(text)
    chat_id = str(row.get("chat_id"))
    key = (chat_id, int(row.get("message_id", -1)))
    name = str(row.get("first_name") or "")
    username = str(row.get("username") or "")
    if not name and not username:
        name, username = names.get(key, ("", ""))
    msg = {
        "text": text,
        "has_link": bool(LINK_RE.search(text)),
        "has_mention": "@" in text,
        "has_media": row.get("kind") in {"image", "sticker", "video", "animation", "video_note"}
                     or text.startswith(("［图片］", "［贴纸］")),
        "is_forwarded": "［转发自］" in text,
        "length": len(text),
    }
    sender = {
        "user_id": row.get("user_id"),
        "username": username,
        "first_name": name,
        "last_name": "",
        "bio": "",
        "bio_links": [],
        "age_hours": 0,
        "age_known": False,
        "msgs_in_group": 0,
        "prior_ad_hits": 0,
    }
    recent = [{"user": str(x.get("user_id", "")),
               "text": str(x.get("text") or x.get("text_raw") or "")[:200]}
              for x in prev_map.get(key, [])]
    return {
        "chat": {"id": row.get("chat_id"), "title": f"群聊 {chat_id}"},
        "message": msg,
        "sender": sender,
        "recent_context": recent,
        "known_ad_patterns": "",
        "known_false_positives": "",
        "matched_rules": [],
    }


class RateLimiter:
    def __init__(self, rpm: float):
        self.interval = 0.0 if rpm <= 0 else 60.0 / max(1.0, rpm)
        self.lock = threading.Lock()
        self.next = 0.0

    def wait(self):
        if self.interval <= 0:
            return
        with self.lock:
            now = time.monotonic()
            delay = max(0.0, self.next - now)
            self.next = max(now, self.next) + self.interval
        if delay:
            time.sleep(delay)


def request_teacher(url, key, model, questions, state, retries=20):
    payload = json.dumps({"model": model, "questions": questions, "state": state},
                         ensure_ascii=False).encode("utf-8")
    last = None
    for attempt in range(retries):
        req = urllib.request.Request(url, data=payload, method="POST", headers={
            "Authorization": "Bearer " + key, "Content-Type": "application/json",
            "Accept": "application/json", "Accept-Encoding": "identity",
            "User-Agent": "laya-ad-cleaner/1.0"})
        try:
            with urllib.request.urlopen(req, timeout=180) as r:
                return json.loads(r.read().decode("utf-8", "replace"))
        except urllib.error.HTTPError as e:
            err_body = e.read().decode('utf-8','replace')
            last = f"HTTP {e.code}: {err_body[:300]}"
            if e.code == 429:
                # 按要求：429 立即重试，不退避（全局 RPM 限速仍生效）
                continue
            if e.code in (500, 502, 503, 504):
                time.sleep(2 * (attempt + 1))
                continue
            break
        except Exception as e:
            last = f"{type(e).__name__}: {e}"
            time.sleep(2 * (attempt + 1))
    raise RuntimeError(last)


def answer_summary(resp):
    a = resp.get("answers") or {}
    out = {}
    for qid in ("is_ad", "ad_kind", "ad_scope", "severity"):
        v = a.get(qid) or {}
        if qid == "severity":
            out[qid] = {"score": v.get("score"), "probabilities": v.get("probabilities"),
                        "confidence": v.get("confidence")}
        else:
            out[qid] = {"choice": v.get("choice"), "probabilities": v.get("probabilities"),
                        "confidence": v.get("confidence")}
    return out, resp.get("usage") or {}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="data/processed/archive/messages_labeled.jsonl")
    ap.add_argument("--template", default=".work/teacher/jev_request_template.json")
    ap.add_argument("--out", default="data/cleaned/teacher_sample.jsonl")
    ap.add_argument("--per_class", type=int, default=100)
    ap.add_argument("--workers", type=int, default=256)
    ap.add_argument("--rpm", type=float, default=0, help="global requests per minute; 0 = unlimited")
    ap.add_argument("--all", action="store_true", help="clean all candidates (no sampling)")
    ap.add_argument("--resume", action="store_true", help="skip rows already present in --out")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--include_all", action="store_true", help="include all deciders, not only systemone")
    args = ap.parse_args()

    cfg = load_env()
    template = json.load(open(args.template, encoding="utf-8"))
    questions = template["questions"]

    rows = []
    all_rows = load_jsonl(args.data)
    for r in all_rows:
        if not args.include_all and r.get("decider") != "systemone":
            continue
        if r.get("label") not in ("ad", "clean"):
            continue
        if not (r.get("text") or r.get("text_raw")):
            continue
        rows.append(r)
    print("candidates:", len(rows), Counter(r["label"] for r in rows))

    rng = random.Random(args.seed)
    by = defaultdict(list)
    for r in rows:
        by[r["label"]].append(r)
    sample = []
    if args.all:
        sample = list(rows)
        rng.shuffle(sample)
    else:
        for lab in ("ad", "clean"):
            pool = by.get(lab, [])
            rng.shuffle(pool)
            sample.extend(pool[:args.per_class])
        rng.shuffle(sample)
    print("sample:", len(sample), Counter(r["label"] for r in sample))

    if args.resume and Path(args.out).exists():
        done = set()
        for line in open(args.out, encoding="utf-8"):
            try:
                r = json.loads(line)
                done.add((str(r.get("chat_id")), int(r.get("message_id", -1))))
            except Exception:
                continue
        before = len(sample)
        sample = [r for r in sample if (str(r.get("chat_id")), int(r.get("message_id", -1))) not in done]
        print(f"resume: skip {before-len(sample)}, remaining {len(sample)}")

    # 同群前序消息（用于 recent_context；用全部带文本消息，不只 systemone）
    ordered = sorted(
        (r for r in all_rows if (r.get("text") or r.get("text_raw"))),
        key=lambda r: (str(r.get("chat_id")), int(r.get("date", 0)), int(r.get("message_id", 0))))
    prev_map = {}
    for cid, group in groupby(ordered, key=lambda r: str(r.get("chat_id"))):
        g = list(group)
        for i, r in enumerate(g):
            prev_map[(cid, int(r.get("message_id", -1)))] = g[max(0, i - 2):i]
    names = load_judgment_names()
    limiter = RateLimiter(args.rpm)

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    lock = threading.Lock()
    results = []
    t0 = time.time()

    def work(r):
        state = make_state(r, names, prev_map)
        limiter.wait()
        resp = request_teacher(cfg["TEACHER_URL"], cfg["TEACHER_KEY"],
                               cfg.get("TEACHER_MODEL", "jev-1.13"), questions, state)
        summary, usage = answer_summary(resp)
        return {"chat_id": r.get("chat_id"), "message_id": r.get("message_id"),
                "text": r.get("text") or r.get("text_raw"),
                "v1_label": r.get("label"), "v1_kind": r.get("ad_kind"),
                "v1_confidence": r.get("confidence"), "split": r.get("split"),
                "v1_decider": r.get("decider"), "v1_scope": r.get("v1_scope"),
                "v1_severity": r.get("v1_severity"),
                "teacher": summary, "usage": usage, "model": resp.get("model")}

    errors = []
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(work, r): r for r in sample}
        done = 0
        for fut in as_completed(futs):
            try:
                res = fut.result()
                with lock:
                    results.append(res)
                    with open(args.out, "a", encoding="utf-8") as f:
                        f.write(json.dumps(res, ensure_ascii=False) + "\n")
            except Exception as e:
                errors.append({"error": str(e), "row": futs[fut].get("message_id")})
            done += 1
            if done % 10 == 0:
                print(f"progress {done}/{len(sample)} elapsed {time.time()-t0:.0f}s", flush=True)

    # 报告
    def tlab(r):
        return (r["teacher"].get("is_ad") or {}).get("choice")
    agree = sum(1 for r in results if (r["v1_label"] == "ad") == (tlab(r) == "ad"))
    conf = Counter((r["v1_label"], tlab(r)) for r in results)
    kinds = Counter((r["teacher"].get("ad_kind") or {}).get("choice") for r in results)
    scopes = Counter((r["teacher"].get("ad_scope") or {}).get("choice") for r in results)
    sevs = Counter((r["teacher"].get("severity") or {}).get("score") for r in results)
    v1_ad_kind = Counter(r["v1_kind"] for r in results if r["v1_label"] == "ad")
    by_dec = Counter(r.get("v1_decider") for r in results)
    agree_dec = {}
    for dec in by_dec:
        sub = [r for r in results if r.get("v1_decider") == dec]
        agree_dec[dec] = round(sum(1 for r in sub if (r["v1_label"] == "ad") == (tlab(r) == "ad")) / len(sub), 4)
    report = {
        "n": len(results), "errors": len(errors),
        "agree_ad_clean": round(agree / max(1, len(results)), 4),
        "confusion": {f"v1_{a}->teacher_{b}": c for (a, b), c in conf.items()},
        "teacher_kind": dict(kinds), "teacher_scope": dict(scopes),
        "teacher_severity_score": {str(k): v for k, v in sevs.items()},
        "v1_ad_kind": dict(v1_ad_kind),
        "by_decider": dict(by_dec), "agree_by_decider": agree_dec,
        "v1_scope_ad": dict(Counter(r.get("v1_scope") for r in results if r["v1_label"] == "ad")),
        "v1_severity_ad": dict(Counter(str(r.get("v1_severity")) for r in results if r["v1_label"] == "ad")),
        "elapsed_s": round(time.time() - t0, 1),
        "errors_detail": errors[:20],
    }
    out = Path(args.out)
    (out.parent / (out.stem + "_report.json")).write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    md = [
        f"# jev-1.13 清洗样本报告\n",
        f"- 样本 {report['n']} 条（v1 ad/clean 各 {args.per_class}），错误 {report['errors']}",
        f"- v1 与 teacher ad/clean 一致率：**{report['agree_ad_clean']:.2%}**",
        f"- 混淆：{report['confusion']}",
        f"- teacher kind：{report['teacher_kind']}",
        f"- teacher scope：{report['teacher_scope']}",
        f"- teacher severity score：{report['teacher_severity_score']}",
        f"- 按 decider 一致率：{report['agree_by_decider']}",
        f"- v1 广告的 scope：{report['v1_scope_ad']}",
        f"- v1 广告的 severity：{report['v1_severity_ad']}",
        f"- 耗时：{report['elapsed_s']}s（workers={args.workers}）",
    ]
    (out.parent / (out.stem + "_report.md")).write_text("\n".join(md), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))
    print("->", out.resolve())


if __name__ == "__main__":
    main()