#!/usr/bin/env python3
"""构建广告细类（7 类）数据集，复用 v1.1 的 train/val split，避免泄漏。

类别：crypto / porn / porn_bait / gambling / scam / promo / spam_flood
映射：contact → porn（样本均为色情招揽/迷药）；spam / impersonate / unspecified 丢弃（样本极少且语义杂）
测试集：
  gold_val     = processed_v2/val ∩ gold.kind（广告行）
  archive_test = processed_v2/val ∩ archive.strong（ad 行）
训练池：
  processed_v2/train ∩ (gold.kind 广告行 ∪ archive.strong ad 行)，85/15 分层切分 train/val
"""
import argparse
import collections
import json
from pathlib import Path

import numpy as np
from sklearn.model_selection import train_test_split

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]
MAP = {"contact": "porn"}
DROP = {"spam", "impersonate", "unspecified", "normal", "none", None, ""}


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


def norm(t):
    return t.strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold", default="data/raw/dataset/converted/gold.kind.jsonl")
    ap.add_argument("--archive", default="data/raw/archive/converted/archive.strong.binary.jsonl")
    ap.add_argument("--split_train", default="data/processed_v2/train.jsonl")
    ap.add_argument("--split_val", default="data/processed_v2/val.jsonl")
    ap.add_argument("--out", default="data/processed_kind")
    ap.add_argument("--val_ratio", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

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

    gold_kind = {}
    for r in load(args.gold):
        lab = clean(r["label"])
        if lab:
            gold_kind.setdefault(norm(r["text"]), lab)
    arch_kind = {}
    for r in load(args.archive):
        if r.get("label") != "ad":
            continue
        lab = clean(r.get("ad_kind"))
        if lab:
            arch_kind.setdefault(norm(r["text"]), lab)

    tr_texts = {norm(r["text"]) for r in load(args.split_train)}
    va_texts = {norm(r["text"]) for r in load(args.split_val)}

    pool = ([{"text": t, "label": lab} for t, lab in gold_kind.items() if t in tr_texts] +
            [{"text": t, "label": lab} for t, lab in arch_kind.items()
             if t in tr_texts and t not in gold_kind])
    gold_val = [{"text": t, "label": lab} for t, lab in gold_kind.items() if t in va_texts]
    archive_test = [{"text": t, "label": lab} for t, lab in arch_kind.items()
                    if t in va_texts and t not in gold_kind]

    assert not ({norm(r["text"]) for r in pool} & {norm(r["text"]) for r in gold_val})
    assert not ({norm(r["text"]) for r in pool} & {norm(r["text"]) for r in archive_test})

    labels = np.array([r["label"] for r in pool])
    idx = np.arange(len(pool))
    tr_idx, va_idx = train_test_split(idx, test_size=args.val_ratio,
                                      stratify=labels, random_state=args.seed)
    train = [pool[i] for i in tr_idx]
    val = [pool[i] for i in va_idx]

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in [("train", train), ("val", val),
                       ("gold_val", gold_val), ("archive_test", archive_test)]:
        with open(out / f"{name}.jsonl", "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"{name:<14} n={len(rows):5d} {dict(collections.Counter(r['label'] for r in rows))}")
    label_map = {k: i for i, k in enumerate(sorted(KINDS))}
    (out / "label_map.json").write_text(
        json.dumps(label_map, ensure_ascii=False, indent=2), encoding="utf-8")
    (out / "report.json").write_text(json.dumps({
        "kinds": KINDS, "label_map": label_map,
        "train": len(train), "val": len(val),
        "gold_val": len(gold_val), "archive_test": len(archive_test),
        "train_dist": dict(collections.Counter(r["label"] for r in train)),
        "val_dist": dict(collections.Counter(r["label"] for r in val)),
        "gold_val_dist": dict(collections.Counter(r["label"] for r in gold_val)),
        "archive_test_dist": dict(collections.Counter(r["label"] for r in archive_test)),
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print("输出:", out.resolve())


if __name__ == "__main__":
    main()