#!/usr/bin/env python3
"""从 archive.tar.gz 解包出的群聊存档构建清洗后的带标签数据集。

输入目录结构（bot 存档导出）：
    <archive>/<chat_id>/messages-YYYY-MM-DD.jsonl   消息
    <archive>/<chat_id>/judgments-YYYY-MM-DD.jsonl  判定（含 verdict/ad_kind）
    <archive>/<chat_id>/images/YYYY-MM-DD/<msgid>_{tg|hist_*}.{jpg,webp}

三份产物：
    messages.jsonl   全量消息 + 裁决后的标签（ad/clean）
    image_clf.jsonl  通过校验的图片 + 图文对齐后的标签
    desc_clf.jsonl   图片/贴纸的 VLM 描述文本 + 标签（喂文本模型用）

清洗口径见 report.md。原始存档只读，产物全部写到 --out。
"""

from __future__ import annotations

import argparse
import collections
import datetime as dt
import hashlib
import json
import os
import re
import sys
import unicodedata
from typing import Any, Iterable

# 判定仲裁优先级：数字越小越可信
TIER_NAMES = {
    "manual": 0, "manual-ban": 0,
    "systemone+llm": 1,
    "systemone": 2,
    "llm": 3, "llm+llm": 3,
    "phash": 5, "hash": 5, "hash+llm": 5,
}
DROP_VERDICTS = {"skipped", "error"}
# 反检测用的不可见字符：零宽空格/连接符、词连接符、BOM、软连字符
INVISIBLE = dict.fromkeys(map(ord, "\u200b\u200c\u200d\u2060\u2061\u2062\u2063\u2064\ufeff\u00ad"), None)

MARK_IMAGE = "［图片］"
MARK_STICKER = "［贴纸］"
MARK_QUOTE = "［引用］"

# 群机器人指令/签到类短文本：样本量大但不含判定信息
COMMAND_WORDS = {"签到", "专享", "订阅", "积分", "积分排名", "积分兑换", "1", "?", ".", "。", "+1"}

# 联系方式/落地页特征：短文本里带这些字样的 ad 才算内容级广告
CONTACT_RE = re.compile(
    r"(https?://|t\.me/|www\.|\+\d{5,}|\d{3,}|@[A-Za-z0-9_]{4,}|微信|威信|vx|VX|QQ|qq|telegram|电报|飞机|私聊|加我)",
    re.IGNORECASE,
)


def label_trust(label: str | None, decider: str | None, text: str) -> str:
    """标注可信度分层。

    存档里的判定不全是内容级：manual 与"极短文本被判 ad"多数是发送者级处置
    （用户因历史行为被封，其所有消息都被打成 ad），拿来训练会教模型把"嗯""哈喽"学成广告。
    """
    if label not in ("ad", "clean"):
        return "none"
    if decider == "manual":
        return "suspect_sender_level"
    if label == "ad" and len(text) < 6 and not CONTACT_RE.search(text):
        return "suspect_sender_level"
    if decider and (decider.startswith("rule") or decider in ("phash", "hash", "hash+llm")):
        return "medium_heuristic"
    return "high"

IMG_RE = re.compile(r"^(\d+)_(tg|hist|\d+)[^/]*\.(jpg|jpeg|png|webp)$", re.IGNORECASE)


def iter_jsonl(pattern_dir: str, prefix: str) -> Iterable[dict]:
    for root, _dirs, files in os.walk(pattern_dir):
        for name in files:
            if not (name.startswith(prefix) and name.endswith(".jsonl")):
                continue
            path = os.path.join(root, name)
            with open(path, encoding="utf-8") as fh:
                for lineno, line in enumerate(fh, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        yield json.loads(line)
                    except json.JSONDecodeError:
                        print(f"[warn] 坏行 {path}:{lineno}", file=sys.stderr)


def tier_of(decider: str | None) -> int:
    if decider is None:
        return 9
    if decider in TIER_NAMES:
        return TIER_NAMES[decider]
    if decider.startswith("rule"):
        return 4
    return 6


def load_messages(base: str) -> dict[tuple[str, int], dict]:
    """按 (chat_id, message_id) 去重，保留 edit_date 最新的一条。"""
    msgs: dict[tuple[str, int], dict] = {}
    total = dup = 0
    for d in iter_jsonl(base, "messages-"):
        total += 1
        key = (str(d.get("chat_id")), int(d.get("message_id", -1)))
        prev = msgs.get(key)
        if prev is None:
            msgs[key] = d
            continue
        dup += 1
        if int(d.get("edit_date") or 0) >= int(prev.get("edit_date") or 0):
            msgs[key] = d
    print(f"[load] messages 原始 {total} 条，去重后 {len(msgs)}（重复 {dup}）")
    return msgs


def load_judgments(base: str) -> dict[tuple[str, int], list[dict]]:
    groups: dict[tuple[str, int], list[dict]] = collections.defaultdict(list)
    total = 0
    for d in iter_jsonl(base, "judgments-"):
        total += 1
        groups[(str(d.get("chat_id")), int(d.get("message_id", -1)))].append(d)
    print(f"[load] judgments {total} 条，覆盖 {len(groups)} 个消息")
    return dict(groups)


def resolve_label(js: list[dict] | None) -> dict[str, Any]:
    """多判定仲裁：先按 decider 可信度分层，层内按置信度、再按时间取最新。

    同层出现 ad/clean 对立即记为冲突（保留高置信度那条，但标记出来供人工复核）。
    """
    empty = {"label": None, "ad_kind": None, "confidence": None, "decider": None,
             "label_tier": None, "label_conflict": False, "label_votes": 0}
    if not js:
        return empty
    valid = [j for j in js if j.get("verdict") not in DROP_VERDICTS and j.get("verdict")]
    if not valid:
        return empty
    best_tier = min(tier_of(j.get("decider")) for j in valid)
    top = [j for j in valid if tier_of(j.get("decider")) == best_tier]
    votes = {j["verdict"] for j in top}
    top.sort(key=lambda j: (float(j.get("confidence") or 0), int(j.get("created_at") or 0)), reverse=True)
    win = top[0]
    return {
        "label": win["verdict"],
        "ad_kind": win.get("ad_kind"),
        "confidence": win.get("confidence"),
        "decider": win.get("decider"),
        "label_tier": best_tier,
        "label_conflict": len(votes) > 1,
        "label_votes": len(valid),
    }


def clean_text(raw: str) -> dict[str, Any]:
    """剥离不可见字符、拆出引用段、判定内容类型。

    标记（［图片］等）必须在 NFKC 归一化前识别，否则全角括号会被改写成半角。
    """
    raw = raw or ""
    had_invisible = any(ch in raw for ch in "\u200b\u200c\u200d\u2060\u2061\u2062\u2063\u2064\ufeff\u00ad")
    content_type = "text"
    if raw.startswith(MARK_IMAGE):
        content_type = "image_desc"
    elif raw.startswith(MARK_STICKER):
        content_type = "sticker_desc"

    body, quote = raw, ""
    if MARK_QUOTE in raw:
        head, _, tail = raw.partition(MARK_QUOTE)
        body, quote = (head if head.strip() else raw), tail.strip()

    def norm(s: str) -> str:
        s = unicodedata.normalize("NFKC", s)
        s = s.translate(INVISIBLE)
        return re.sub(r"[ \t\u3000]+", " ", re.sub(r"\n{3,}", "\n\n", s)).strip()

    return {
        "text": norm(body),
        "quote": norm(quote),
        "content_type": content_type,
        "had_invisible": had_invisible,
        "had_quote": bool(quote),
    }


def dhash(path: str, size: int = 8) -> int | None:
    """8x8 差值哈希，用于跨消息的近似重复图去重。"""
    from PIL import Image
    try:
        with Image.open(path) as im:
            im = im.convert("L").resize((size + 1, size), Image.Resampling.LANCZOS)
            px = list(im.getdata())
    except Exception:
        return None
    bits = 0
    for row in range(size):
        for col in range(size):
            left = px[row * (size + 1) + col]
            right = px[row * (size + 1) + col + 1]
            bits = (bits << 1) | (1 if left > right else 0)
    return bits


def hamming(a: int, b: int) -> int:
    return (a ^ b).bit_count()


def collect_images(base: str) -> tuple[dict[tuple[str, int], list[dict]], dict[str, int]]:
    """枚举图片文件，校验可解码，按消息归组。"""
    from PIL import Image

    stats = collections.Counter()
    per_msg: dict[tuple[str, int], list[dict]] = collections.defaultdict(list)
    for root, _dirs, files in os.walk(base):
        if os.path.basename(os.path.dirname(root)) != "images":
            continue
        chat_id = os.path.basename(os.path.dirname(os.path.dirname(root)))
        for name in files:
            stats["files"] += 1
            m = IMG_RE.match(name)
            if not m:
                stats["unparsable_name"] += 1
                continue
            path = os.path.join(root, name)
            partial = ".partial." in name.lower()
            if partial:
                stats["partial"] += 1
            try:
                with Image.open(path) as im:
                    im.load()
                    w, h = im.size
                    fmt = im.format
            except Exception:
                stats["corrupt"] += 1
                continue
            if w < 64 or h < 64:
                stats["too_small"] += 1
                continue
            with open(path, "rb") as fh:
                digest = hashlib.sha256(fh.read()).hexdigest()
            stats["valid"] += 1
            per_msg[(chat_id, int(m.group(1)))].append({
                "path": path, "sha256": digest, "w": w, "h": h, "format": fmt,
                "partial": partial, "variant": m.group(2), "bytes": os.path.getsize(path),
            })
    print(f"[img] {dict(stats)}")
    return per_msg, dict(stats)


def pick_one(cands: list[dict]) -> dict:
    """同一消息多份文件时择优：非 partial > 面积大 > 体积大。"""
    return sorted(cands, key=lambda c: (c["partial"], -c["w"] * c["h"], -c["bytes"]))[0]


def build_splits(records: list[dict], val_ratio: float = 0.1, group_keys=None) -> None:
    """时间切分：按 date 排序，尾部 val_ratio 进 val。

    同组记录强制同侧（同一文本 / 同一相册），避免近似样本跨集泄漏。
    三份数据集各自按自身时间分布切分，否则图片这类集中在新日期的子集会全掉进 val。
    """
    dates = sorted(r["date"] for r in records)
    if not dates:
        return
    cutoff = dates[int(len(dates) * (1 - val_ratio))]
    for r in records:
        r["split"] = "val" if r["date"] >= cutoff else "train"
    if group_keys is None:
        return

    anchor: dict[str, str] = {}
    for r in sorted(records, key=lambda x: x["date"]):
        keys = group_keys(r)
        first = next((anchor[k] for k in keys if k in anchor), None)
        if first is None:
            for k in keys:
                anchor[k] = r["split"]
        else:
            r["split"] = first
            for k in keys:
                anchor[k] = first


def msg_group_keys(r: dict) -> list[str]:
    keys = [f"t:{r['text_hash']}"] if r.get("text") else []
    if r.get("media_group_id") is not None:
        keys.append(f"g:{r['chat_id']}:{r['media_group_id']}")
    return keys


def desc_group_keys(r: dict) -> list[str]:
    return [f"t:{hashlib.sha1(r['text'].encode()).hexdigest()[:16]}"] if r.get("text") else []


def write_jsonl(path: str, rows: Iterable[dict]) -> int:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    n = 0
    with open(path, "w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            n += 1
    return n


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--archive", default="data/raw/archive/archive")
    ap.add_argument("--out", default="data/processed/archive")
    ap.add_argument("--val-ratio", type=float, default=0.1)
    ap.add_argument("--phash-radius", type=int, default=4, help="跨消息近似重复图的汉明距离阈值")
    args = ap.parse_args()

    msgs = load_messages(args.archive)
    judgments = load_judgments(args.archive)

    # ---------- 消息：清洗 + 打标 ----------
    rows: list[dict] = []
    text_seen: dict[str, tuple[str, int]] = {}
    conflicts: list[dict] = []
    for key, m in msgs.items():
        chat_id, mid = key
        meta = clean_text(m.get("text") or "")
        lab = resolve_label(judgments.get(key))
        text_hash = hashlib.sha1(meta["text"].encode("utf-8")).hexdigest()[:16]
        # 首次出现决定该文本的 split 锚点；重复文本标记出来便于降采样
        is_dup = text_hash in text_seen and bool(meta["text"])
        if meta["text"] and not is_dup:
            text_seen[text_hash] = ("train", 0)
        if lab["label_conflict"]:
            conflicts.append({"chat_id": chat_id, "message_id": mid, "text": meta["text"][:120]})
        rows.append({
            "chat_id": chat_id,
            "message_id": mid,
            "date": int(m.get("date") or 0),
            "iso_time": dt.datetime.fromtimestamp(int(m.get("date") or 0)).isoformat(timespec="seconds"),
            "user_id": m.get("user_id"),
            "kind": m.get("kind"),
            "media_group_id": m.get("media_group_id"),
            "punished": bool(m.get("punished")),
            "text_raw": m.get("text") or "",
            "text": meta["text"],
            "quote": meta["quote"],
            "content_type": meta["content_type"],
            "had_invisible": meta["had_invisible"],
            "text_len": len(meta["text"]),
            "low_info": len(meta["text"]) < 4 or meta["text"] in COMMAND_WORDS,
            "text_hash": text_hash,
            "dup_text": is_dup,
            "label_trust": label_trust(lab["label"], lab["decider"], meta["text"]),
            **lab,
        })

    labeled = [r for r in rows if r["label"] in ("ad", "clean")]
    build_splits(labeled, args.val_ratio, msg_group_keys)

    # ---------- 图片：校验 + 去重 + 对齐标签 ----------
    per_msg, img_stats = collect_images(args.archive)
    by_key = {(r["chat_id"], r["message_id"]): r for r in rows}
    image_rows: list[dict] = []
    seen_hash: dict[str, dict] = {}
    seen_phash: list[tuple[int, dict]] = []
    img_dedup = collections.Counter()

    for key, cands in sorted(per_msg.items(), key=lambda kv: (kv[0][0], kv[0][1])):
        rec = by_key.get(key)
        if rec is None:
            img_dedup["no_message"] += 1
            continue
        if rec["label"] not in ("ad", "clean"):
            img_dedup["no_label"] += 1
            continue
        best = pick_one(cands)
        if best["sha256"] in seen_hash:
            img_dedup["exact_dup"] += 1
            continue
        ph = dhash(best["path"])
        if ph is not None and any(hamming(ph, p) <= args.phash_radius for p, _ in seen_phash):
            img_dedup["near_dup"] += 1
            continue
        out = {
            "path": os.path.relpath(best["path"], args.archive).replace("\\", "/"),
            "sha256": best["sha256"],
            "dhash": f"{ph:016x}" if ph is not None else None,
            "width": best["w"], "height": best["h"], "bytes": best["bytes"],
            "chat_id": key[0], "message_id": key[1],
            "date": rec["date"], "split": rec["split"],
            "label": rec["label"], "ad_kind": rec["ad_kind"],
            "confidence": rec["confidence"], "decider": rec["decider"],
            "content_type": rec["content_type"], "text": rec["text"],
            "variants": len(cands),
        }
        image_rows.append(out)
        seen_hash[best["sha256"]] = out
        if ph is not None:
            seen_phash.append((ph, out))

    # 图片按自身时间分布切分：它们集中在近期，跟着消息切分会让 val 反超 train
    build_splits(image_rows, args.val_ratio)

    # ---------- 描述文本（图片/贴纸的 VLM 描述）：喂文本模型的桥接训练集 ----------
    img_keys = {(r["chat_id"], r["message_id"]) for r in image_rows}
    desc_rows = []
    for r in rows:
        if r["content_type"] not in ("image_desc", "sticker_desc"):
            continue
        if r["label"] not in ("ad", "clean"):
            continue
        desc_rows.append({
            "chat_id": r["chat_id"], "message_id": r["message_id"], "date": r["date"],
            "split": r["split"], "text": r["text"], "quote": r["quote"],
            "content_type": r["content_type"], "label": r["label"],
            "ad_kind": r["ad_kind"], "confidence": r["confidence"],
            "decider": r["decider"], "has_image": (r["chat_id"], r["message_id"]) in img_keys,
        })
    build_splits(desc_rows, args.val_ratio, desc_group_keys)

    # ---------- 落盘 + 报告 ----------
    out = args.out
    n_msg = write_jsonl(os.path.join(out, "messages.jsonl"), rows)
    n_img = write_jsonl(os.path.join(out, "image_clf.jsonl"), image_rows)
    n_desc = write_jsonl(os.path.join(out, "desc_clf.jsonl"), desc_rows)
    n_lab = write_jsonl(os.path.join(out, "messages_labeled.jsonl"),
                        (r for r in rows if r["label"] in ("ad", "clean")))
    # 去重视图：同一文本只保留首条（"签到"这类群指令单条就能刷到 7420 次）
    seen_h: set[str] = set()
    dedup_rows = []
    for r in labeled:
        if r["text"] and r["text_hash"] in seen_h:
            continue
        seen_h.add(r["text_hash"])
        dedup_rows.append(r)
    n_dedup = write_jsonl(os.path.join(out, "messages_labeled_dedup.jsonl"), dedup_rows)
    # 可信子集：剔除发送者级处置造成的伪标签，供训练直接用
    trusted = [r for r in rows
               if r["label"] in ("ad", "clean") and r["label_trust"] in ("high", "medium_heuristic")]
    n_trusted = write_jsonl(os.path.join(out, "messages_trusted.jsonl"), trusted)

    report = {
        "generated_at": dt.datetime.now().isoformat(timespec="seconds"),
        "archive": os.path.abspath(args.archive),
        "messages": {
            "total": n_msg,
            "labeled": n_lab,
            "label_dist": dict(collections.Counter(r["label"] for r in rows)),
            "ad_kind_dist": dict(collections.Counter(r["ad_kind"] for r in rows if r["label"] == "ad").most_common()),
            "content_type_dist": dict(collections.Counter(r["content_type"] for r in rows)),
            "kind_dist": dict(collections.Counter(r["kind"] for r in rows)),
            "labeled_content_type_dist": dict(collections.Counter(r["content_type"] for r in labeled)),
            "low_info": sum(1 for r in labeled if r["low_info"]),
            "label_trust": dict(collections.Counter(r["label_trust"] for r in labeled)),
            "trusted": {
                "rows": n_trusted,
                "label_dist": dict(collections.Counter(r["label"] for r in trusted)),
                "split": dict(collections.Counter(r["split"] for r in trusted)),
            },
            "labeled_dedup": {
                "rows": n_dedup,
                "label_dist": dict(collections.Counter(r["label"] for r in dedup_rows)),
                "split": dict(collections.Counter(r["split"] for r in dedup_rows)),
            },
            "conflicts": len(conflicts),
            "had_invisible": sum(1 for r in rows if r["had_invisible"]),
            "dup_text": sum(1 for r in rows if r["dup_text"]),
            "split": dict(collections.Counter(r["split"] for r in labeled)),
            "date_range": [dt.datetime.fromtimestamp(labeled[0]["date"]).isoformat(timespec="seconds"),
                           dt.datetime.fromtimestamp(labeled[-1]["date"]).isoformat(timespec="seconds")] if labeled else None,
        },
        "images": {
            "files_scanned": img_stats.get("files", 0),
            "invalid_or_dropped": {k: img_stats.get(k, 0) for k in ("corrupt", "partial", "too_small", "unparsable_name")},
            "dropped_in_join": dict(img_dedup),
            "kept": n_img,
            "label_dist": dict(collections.Counter(r["label"] for r in image_rows)),
            "content_type_dist": dict(collections.Counter(r["content_type"] for r in image_rows)),
            "ad_kind_dist": dict(collections.Counter(r["ad_kind"] for r in image_rows if r["label"] == "ad").most_common()),
            "split": dict(collections.Counter(r["split"] for r in image_rows)),
        },
        "desc": {
            "kept": n_desc,
            "label_dist": dict(collections.Counter(r["label"] for r in desc_rows)),
            "with_image_file": sum(1 for r in desc_rows if r["has_image"]),
            "split": dict(collections.Counter(r["split"] for r in desc_rows)),
        },
        "conflict_samples": conflicts[:20],
    }
    with open(os.path.join(out, "report.json"), "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)

    m, im, de = report["messages"], report["images"], report["desc"]
    md = [
        "# archive 存档清洗报告", "",
        f"- 生成时间：{report['generated_at']}", f"- 源目录：`{report['archive']}`", "",
        "## 消息", "",
        f"- 总消息 {m['total']}，其中带标签 {m['labeled']}",
        f"- 标签分布：`{m['label_dist']}`",
        f"- 已标内容类型：`{m['labeled_content_type_dist']}`",
        f"- 判定冲突（同层 ad/clean 对立）：{m['conflicts']}",
        f"- 含不可见字符（反检测）：{m['had_invisible']}",
        f"- 完全重复文本：{m['dup_text']}；低信息短文本：{m['low_info']}",
        f"- 切分：`{m['split']}`",
        f"- 标注可信度：`{m['label_trust']}`",
        f"- 可信训练集 `messages_trusted.jsonl`：{m['trusted']['rows']} 条，"
        f"标签 `{m['trusted']['label_dist']}`，切分 `{m['trusted']['split']}`",
        f"- 去重视图 `messages_labeled_dedup.jsonl`：{m['labeled_dedup']['rows']} 条，"
        f"标签 `{m['labeled_dedup']['label_dist']}`", "",
        "## 图片", "",
        f"- 扫描 {im['files_scanned']} 个文件，丢弃 `{im['invalid_or_dropped']}`，join 丢弃 `{im['dropped_in_join']}`",
        f"- 保留 {im['kept']}，标签分布 `{im['label_dist']}`，内容类型 `{im['content_type_dist']}`",
        f"- 切分：`{im['split']}`", "",
        "## 图片描述文本", "",
        f"- 保留 {de['kept']}，其中 {de['with_image_file']} 条有对应图片文件",
        f"- 标签分布：`{de['label_dist']}`",
    ]
    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"\n[out] {os.path.abspath(out)}")


if __name__ == "__main__":
    main()