#!/usr/bin/env python3
"""合并"现有训练集"与"新增语料"成一份微调用数据集。

新语料在构建时已排除与现有语料的重复，所以这里只需拼接 + 统一标签口径 + 重算切分。
现有 val 保持独立（它代表生产分布），新语料的 val 单独留作域外评测集。
"""

from __future__ import annotations

import argparse
import collections
import json
import os
import re
import unicodedata


def norm(t: str) -> str:
    t = unicodedata.normalize("NFKC", t).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)
    return re.sub(r"\s+", " ", t).strip()


def load(path: str) -> list[dict]:
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-train", default="data/processed/train.jsonl")
    ap.add_argument("--base-val", default="data/processed/val.jsonl")
    ap.add_argument("--new-train", default="data/processed/text_corpus/train.jsonl")
    ap.add_argument("--new-val", default="data/processed/text_corpus/val.jsonl")
    ap.add_argument("--out", default="data/processed/merged")
    args = ap.parse_args()

    base_tr = load(args.base_train)
    base_va = load(args.base_val)
    new_tr = load(args.new_train)
    new_va = load(args.new_val)

    # 以现有 val 的归一化键为排除集：新语料绝不能出现在生产 val 里
    holdout = {norm(r["text"]) for r in base_va}

    merged, seen = [], set()
    dropped = collections.Counter()
    for src, rows in (("base", base_va + base_tr), ("new", new_tr)):
        for r in rows:
            t, lab = (r.get("text") or "").strip(), str(r.get("label") or "")
            if not t or lab not in ("ad", "normal"):
                dropped["bad_row"] += 1
                continue
            k = norm(t)
            if src == "new" and k in holdout:
                dropped["in_base_val"] += 1
                continue
            # 只在新语料内部去重；现有语料本身保留原样
            if src == "new":
                if k in seen:
                    dropped["dup_new"] += 1
                    continue
                seen.add(k)
            merged.append({"text": t, "label": lab})

    train = [r for r in merged if norm(r["text"]) not in holdout]
    val = base_va

    os.makedirs(args.out, exist_ok=True)
    for name, data in (("train.jsonl", train), ("val.jsonl", val)):
        with open(os.path.join(args.out, name), "w", encoding="utf-8") as fh:
            for r in data:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(os.path.join(args.out, "label_map.json"), "w", encoding="utf-8") as fh:
        json.dump({"ad": 0, "normal": 1}, fh, indent=2)

    print(json.dumps({
        "train": len(train), "val": len(val),
        "train_dist": dict(collections.Counter(r["label"] for r in train)),
        "val_dist": dict(collections.Counter(r["label"] for r in val)),
        "dropped": dict(dropped),
        "out": os.path.abspath(args.out),
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()