#!/usr/bin/env python3
"""构建四头数据集：二分类 + 7 类 kind + ad_scope + severity。

标签来源：data/raw/dataset/ft/antiad_ft_gold.jsonl.gz 的 llm/so.scope、llm/so.severity。
- scope: account=0, message=1；normal 一律 message；广告缺失 = -100
- severity: 0/1/2/3 四档（<0.5 / <1.5 / <2.5 / >=2.5）；normal=0；广告缺失 = -100
- 文本按 convert_dataset 同样的 quality+time 优先级去重
- 复用 processed_dual 的 train/dev/val/holdout split，避免泄漏
"""
import gzip
import json
from pathlib import Path

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]
QUALITY_RANK = {"llm_reviewed": 5, "manual": 4, "rule": 3, "hash": 2, "so_initial": 1, "other": 0}
MAP = {"contact": "porn"}


def clean_kind(lab):
    if lab in MAP:
        lab = MAP[lab]
    return lab if lab in KINDS else None


def sev_bucket(v):
    if v is None:
        return -100
    v = float(v)
    if v < 0.5:
        return 0
    if v < 1.5:
        return 1
    if v < 2.5:
        return 2
    return 3


def main():
    src = Path("data/raw/dataset/ft/antiad_ft_gold.jsonl.gz")
    best = {}
    with gzip.open(src, "rt", encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            text = (r.get("text") or "").strip()
            if not text:
                continue
            rank = QUALITY_RANK.get(r.get("quality"), 0)
            time = r.get("time") or ""
            prev = best.get(text)
            if prev is None or (rank, time) > (prev[0], prev[1]):
                best[text] = (rank, time, r)
    print("unique texts:", len(best))

    labels = {}
    for text, (_, _, r) in best.items():
        verdict = r.get("verdict")
        kind = r.get("ad_kind")
        llm = r.get("llm") or {}
        so = r.get("so") or {}
        scope = llm.get("scope") or so.get("scope")
        sev = llm.get("severity") if llm.get("severity") is not None else so.get("severity")
        if verdict == "ad" and clean_kind(kind):
            label_bin, label_kind = 0, KINDS.index(clean_kind(kind))
            label_scope = 0 if scope == "account" else (1 if scope == "message" else -100)
            label_sev = sev_bucket(sev)
        else:
            label_bin, label_kind = 1, -100
            label_scope, label_sev = 1, 0
        labels[text] = (label_bin, label_kind, label_scope, label_sev)

    out = Path("data/processed_4h")
    out.mkdir(parents=True, exist_ok=True)
    for split in ["train", "dev", "val", "holdout827"]:
        rows = [json.loads(l) for l in open(f"data/processed_dual/{split}.jsonl", encoding="utf-8")]
        new = []
        for row in rows:
            t = row["text"].strip()
            if t in labels:
                b, k, s, v = labels[t]
            else:
                b, k = row["label_bin"], row["label_kind"]
                s, v = (1, 0) if b == 1 else (-100, -100)
            new.append({"text": row["text"], "label_bin": b, "label_kind": k,
                        "label_scope": s, "label_sev": v})
        with open(out / f"{split}.jsonl", "w", encoding="utf-8") as fh:
            for x in new:
                fh.write(json.dumps(x, ensure_ascii=False) + "\n")
        from collections import Counter
        print(split, len(new), "bin", dict(Counter(x["label_bin"] for x in new)),
              "scope", dict(Counter(x["label_scope"] for x in new)),
              "sev", dict(Counter(x["label_sev"] for x in new)))


if __name__ == "__main__":
    main()