#!/usr/bin/env python3
"""用本地微调的 Laya 模型给 archive 存档补标，并统计与生产判定的一致率。

用途有两个：
  1. 质量校验：在已有生产判定（systemone/llm/manual）的消息上跑一遍，
     算一致率 —— 这衡量的是"自训模型 vs 线上系统"的口径差。
  2. 补标：未判定的消息用模型概率补出来，单独落文件并标记 label_source，
     不与生产判定混在一起使用。

产物：
    messages_scored.jsonl   每条有文本的消息 + p_ad
    agreement.json          按 decider / 群 / 内容类型拆分的一致率
"""

from __future__ import annotations

import argparse
import collections
import json
import os
import sys
import time

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

BATCH = 128
MAX_LEN = 512


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--inp", default="data/processed/archive/messages.jsonl")
    ap.add_argument("--model", default="outputs/ad-clf/best")
    ap.add_argument("--out", default="data/processed/archive")
    ap.add_argument("--max-len", type=int, default=MAX_LEN)
    ap.add_argument("--batch", type=int, default=BATCH)
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForSequenceClassification.from_pretrained(args.model).eval().to(device)
    if device == "cuda":
        model = model.to(torch.bfloat16)
    ad_id = int(model.config.label2id["ad"])

    t0 = time.time()
    total = scored = 0
    rows_meta: list[dict] = []
    probs: list[float] = []

    # 按长度排序后分批，减少 padding 浪费；批次内顺序在写回前还原
    items = []
    with open(args.inp, encoding="utf-8") as fh:
        for line in fh:
            r = json.loads(line)
            total += 1
            if not r.get("text"):
                continue
            items.append(r)
    items.sort(key=lambda r: len(r["text"]))
    print(f"[score] 待打分 {len(items)} 条（总 {total}）")

    with torch.inference_mode():
        for i in range(0, len(items), args.batch):
            chunk = items[i:i + args.batch]
            enc = tok([r["text"] for r in chunk], truncation=True, max_length=args.max_len,
                      padding=True, return_tensors="pt").to(device)
            logits = model(**enc).logits.float()
            p = torch.softmax(logits, dim=-1)[:, ad_id].tolist()
            probs.extend(p)
            rows_meta.extend(chunk)
            scored += len(chunk)
            if scored % 20000 < args.batch:
                print(f"  {scored}/{len(items)}  {time.time() - t0:.0f}s")

    # 用 (chat_id, message_id) 作稳定键：重新读文件会生成新对象，不能靠 id()
    scored_map: dict[tuple[str, int], float] = {
        (r["chat_id"], r["message_id"]): probs[k] for k, r in enumerate(rows_meta)
    }
    meta_by_key = {(r["chat_id"], r["message_id"]): r for r in rows_meta}

    def agreement(rows: list[dict]) -> dict:
        ok = collections.Counter()
        n = collections.Counter()
        for r in rows:
            if r["label"] not in ("ad", "clean"):
                continue
            p_ad = scored_map.get((r["chat_id"], r["message_id"]))
            if p_ad is None:
                continue
            pred = "ad" if p_ad >= 0.5 else "clean"
            key_all = "all"
            n[key_all] += 1
            ok[key_all] += pred == r["label"]
            for key in (f"decider={r.get('decider')}", f"chat={r.get('chat_id')}",
                        f"type={r.get('content_type')}"):
                n[key] += 1
                ok[key] += pred == r["label"]
        return {k: {"n": n[k], "agree": round(ok[k] / n[k], 4)} for k in n if n[k] >= 10}

    out_dir = args.out
    out_path = os.path.join(out_dir, "messages_scored.jsonl")
    pseudo = []
    with open(out_path, "w", encoding="utf-8") as fh:
        for line in open(args.inp, encoding="utf-8"):
            r = json.loads(line)
            p_ad = scored_map.get((r["chat_id"], r["message_id"]))
            if p_ad is not None:
                r["p_ad"] = round(p_ad, 6)
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            if r.get("text") and r["label"] not in ("ad", "clean"):
                r2 = dict(r)
                r2.update({"label": "ad" if p_ad >= 0.5 else "clean",
                           "label_source": "laya_pseudo", "p_ad": round(p_ad, 6),
                           "label_tier": 9, "confidence": round(max(p_ad, 1 - p_ad), 4),
                           "decider": "laya_pseudo", "ad_kind": None})
                pseudo.append(r2)

    pseudo_path = os.path.join(out_dir, "messages_pseudo.jsonl")
    with open(pseudo_path, "w", encoding="utf-8") as fh:
        for r in pseudo:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")

    agree = agreement(list(meta_by_key.values()))
    report = {
        "model": args.model,
        "device": device,
        "scored": scored,
        "pseudo_labeled": len(pseudo),
        "pseudo_label_dist": dict(collections.Counter(r["label"] for r in pseudo)),
        "agreement": agree,
        "elapsed_s": round(time.time() - t0, 1),
    }
    with open(os.path.join(out_dir, "agreement.json"), "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)

    print(json.dumps(report, ensure_ascii=False, indent=2))
    print(f"[out] {out_path}\n[out] {pseudo_path}")


if __name__ == "__main__":
    sys.exit(main())