#!/usr/bin/env python3
"""把清洗后的两源语料整理成可直接微调的文本语料。

原始清洗（build_archive_dataset / build_channel_dataset）解决的是"记录级"问题：
坏行、重复消息、冲突判定、发送者级伪标签。这里解决"语料级"问题，缺一个都会毁掉微调：

1. 标签口径统一 —— 存档用 clean，现有语料用 normal，不统一会把同一类拆成两个类。
2. 与现有训练集去重 —— 实测 23.66% 的群聊文本与 data/processed/train|val 完全重复，
   不排除就会让新 val 直接读到训练过的样本。
3. 模板近似去重 —— 广告文本大量是同一模板换 @用户名/链接/数字，
   归一化后仍有 22.9% 冗余，会放大少数模板的权重并跨集泄漏。
4. 分源时间切分 —— 两源时间分布不同，全局切分会让某个源整体落进 train 或 val。
"""

from __future__ import annotations

import argparse
import collections
import datetime as dt
import hashlib
import json
import os
import re
import sys
import unicodedata

LABEL_MAP = {"ad": "ad", "normal": "normal", "clean": "normal"}


def norm_key(text: str) -> str:
    """归一化去重键：抹掉链接/用户名/数字，让同模板的变体落到一起。"""
    t = unicodedata.normalize("NFKC", text).lower()
    t = re.sub(r"https?://\S+|t\.me/\S+|www\.\S+", "<url>", t)
    t = re.sub(r"@[a-z0-9_]{3,}", "<at>", t)
    t = re.sub(r"\d+", "<n>", t)
    t = re.sub(r"\s+", " ", t).strip()
    return t


def h(text: str) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()


def load_jsonl(path: str) -> list[dict]:
    if not os.path.exists(path):
        print(f"[skip] 缺 {path}", file=sys.stderr)
        return []
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def load_existing(paths: list[str]) -> tuple[set[str], set[str]]:
    exact, normed = set(), set()
    for p in paths:
        for r in load_jsonl(p):
            t = (r.get("text") or "").strip()
            if t:
                exact.add(h(t))
                normed.add(norm_key(t))
    return exact, normed


def collect(args) -> tuple[list[dict], dict]:
    rows: list[dict] = []
    stats = collections.Counter()

    group = load_jsonl("data/processed/archive/messages_trusted.jsonl")
    for r in group:
        if r.get("status") == "pseudo" or r.get("label_source") == "laya_pseudo":
            continue
        rows.append({**r, "source": "group"})
        stats["group_raw"] += 1

    if args.include_pseudo:
        for r in load_jsonl("data/processed/archive/messages_pseudo.jsonl"):
            rows.append({**r, "source": "group_pseudo"})
            stats["pseudo_raw"] += 1

    channel = load_jsonl("data/processed/channel/messages.jsonl")
    for r in channel:
        rows.append({**r, "source": "channel"})
        stats["channel_raw"] += 1

    # TGCN 群导出：标签是本地模型伪标签，只有高置信 normal 才收进训练集。
    # 伪标签 ad 不自动进（自己训练自己只会放大现有偏差），进 needs_review 通道。
    if not args.no_tgcn:
        tgcn = load_jsonl("data/processed/chatexport_2026-10-27/messages.jsonl")
        stats["tgcn_raw"] = len(tgcn)
        for r in tgcn:
            p = r.get("p_ad")
            if r.get("label") == "normal" and p is not None and p <= args.tgcn_normal_max:
                rows.append({**r, "source": "tgcn", "ad_kind": None})
                stats["tgcn_kept"] += 1
            elif r.get("label") == "ad":
                stats["tgcn_ad_to_review"] += 1

    out = []
    for r in rows:
        text = (r.get("text") or "").strip()
        label = LABEL_MAP.get(r.get("label"))
        if not text or label is None:
            stats["drop_empty_or_unlabeled"] += 1
            continue
        if r.get("needs_review"):
            stats["drop_needs_review"] += 1
            continue
        if args.drop_low_info and r.get("low_info"):
            stats["drop_low_info"] += 1
            continue
        out.append({"text": text, "label": label, "source": r["source"],
                    "chat_id": r.get("chat_id"), "message_id": r.get("message_id"),
                    "date": int(r.get("date") or 0),
                    "ad_kind": r.get("ad_kind"), "label_trust": r.get("label_trust", "prior"),
                    "norm": norm_key(text)})
    return out, stats


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/processed/text_corpus")
    ap.add_argument("--val-ratio", type=float, default=0.1)
    ap.add_argument("--neg-multiplier", type=float, default=3.0,
                    help="平衡版本里负样本上限 = 正样本 × 该系数")
    ap.add_argument("--include-pseudo", action="store_true", help="把 Laya 补标样本也算进来")
    ap.add_argument("--no-tgcn", action="store_true", help="不并入 TGCN 群导出的高置信 normal")
    ap.add_argument("--tgcn-normal-max", type=float, default=0.1,
                    help="TGCN 伪标签 normal 的收录阈值（p_ad ≤ 该值）")
    ap.add_argument("--drop-low-info", action="store_true", help="丢掉 <4 字/签到类短文本")
    ap.add_argument("--exclude", nargs="*",
                    default=["data/processed/train.jsonl", "data/processed/val.jsonl"],
                    help="现有语料，用于排除重叠样本")
    args = ap.parse_args()

    rows, stats = collect(args)

    # 1) 排除与现有语料的重叠（近似也算）
    ex_exact, ex_norm = load_existing(args.exclude)
    kept = []
    for r in rows:
        if h(r["text"]) in ex_exact or r["norm"] in ex_norm:
            stats["drop_existing_overlap"] += 1
            continue
        kept.append(r)

    # 2) 归一化去重：同模板只留最长的一条；标签冲突的整组丢弃
    groups: dict[str, list[dict]] = collections.defaultdict(list)
    for r in kept:
        groups[r["norm"]].append(r)
    deduped, conflict_groups = [], 0
    for key, grp in groups.items():
        labels = {x["label"] for x in grp}
        if len(labels) > 1:
            conflict_groups += 1
            stats["drop_norm_conflict"] += len(grp)
            continue
        best = max(grp, key=lambda x: (len(x["text"]), -x["date"]))
        best["dup_group"] = len(grp)
        stats["merged_duplicates"] += len(grp) - 1
        deduped.append(best)

    # 3) 分源时间切分，再按归一化键对齐，避免同模板跨集
    for source in {r["source"] for r in deduped}:
        sub = sorted((r for r in deduped if r["source"] == source), key=lambda r: r["date"])
        if not sub:
            continue
        cutoff = sub[int(len(sub) * (1 - args.val_ratio))]["date"]
        for r in sub:
            r["split"] = "val" if r["date"] >= cutoff else "train"
    anchor: dict[str, str] = {}
    for r in sorted(deduped, key=lambda r: r["date"]):
        r["split"] = anchor.setdefault(r["norm"], r["split"])

    train = [r for r in deduped if r["split"] == "train"]
    val = [r for r in deduped if r["split"] == "val"]

    # 4) 平衡版本：压负样本，保留全部正样本
    pos = sum(1 for r in train if r["label"] == "ad")
    neg_cap = int(pos * args.neg_multiplier)
    negs = [r for r in train if r["label"] == "normal"]
    # 按时间等距抽样，避免只留早期样本
    if len(negs) > neg_cap:
        step = len(negs) / neg_cap
        negs = [negs[int(i * step)] for i in range(neg_cap)]
    balanced = [r for r in train if r["label"] == "ad"] + negs
    balanced.sort(key=lambda r: r["date"])

    out = args.out
    os.makedirs(out, exist_ok=True)
    fields = ("text", "label", "source", "chat_id", "message_id", "date", "ad_kind",
              "label_trust", "dup_group")

    def write(name: str, data: list[dict], only_text_label: bool) -> int:
        with open(os.path.join(out, name), "w", encoding="utf-8") as fh:
            for r in data:
                rec = {"text": r["text"], "label": r["label"]} if only_text_label else \
                      {k: r.get(k) for k in fields if k in r}
                fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
        return len(data)

    n_tr = write("train.jsonl", train, True)
    n_va = write("val.jsonl", val, True)
    n_ba = write("train_balanced.jsonl", balanced, True)
    write("manifest.jsonl", train + val, False)
    with open(os.path.join(out, "label_map.json"), "w", encoding="utf-8") as fh:
        json.dump({name: i for i, name in enumerate(sorted({"ad", "normal"}))},
                  fh, ensure_ascii=False, indent=2)

    def dist(data: list[dict]) -> dict:
        return dict(collections.Counter(r["label"] for r in data))

    def by_source(data: list[dict]) -> dict:
        c = collections.Counter((r["source"], r["label"]) for r in data)
        return {f"{s}/{l}": v for (s, l), v in sorted(c.items())}

    report = {
        "input_stats": dict(stats),
        "norm_conflict_groups": conflict_groups,
        "deduped": len(deduped),
        "train": n_tr, "val": n_va, "train_balanced": n_ba,
        "train_dist": dist(train), "val_dist": dist(val), "balanced_dist": dist(balanced),
        "train_by_source": by_source(train), "val_by_source": by_source(val),
        "date_range": [dt.datetime.fromtimestamp(min(r["date"] for r in deduped)).isoformat(timespec="seconds"),
                       dt.datetime.fromtimestamp(max(r["date"] for r in deduped)).isoformat(timespec="seconds")],
        "out": os.path.abspath(out),
    }
    with open(os.path.join(out, "report.json"), "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)

    md = ["# 文本语料构建报告", "",
          f"- 归一化去重后 **{len(deduped)}** 条（丢弃/合并明细见 `input_stats`）",
          f"- 模板冲突组（同模板不同标签，整组丢弃）：{conflict_groups}", "",
          "| 集合 | 条数 | 标签分布 |", "|---|---|---|",
          f"| train | {n_tr} | `{dist(train)}` |",
          f"| val | {n_va} | `{dist(val)}` |",
          f"| train_balanced | {n_ba} | `{dist(balanced)}` |", "",
          f"- train 来源构成：`{by_source(train)}`",
          f"- val 来源构成：`{by_source(val)}`", "",
          "## 输入侧丢弃统计", "", f"`{dict(stats)}`"]
    with open(os.path.join(out, "report.md"), "w", encoding="utf-8") as fh:
        fh.write("\n".join(md) + "\n")
    print("\n".join(md))
    print(f"[out] {os.path.abspath(out)}")


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