#!/usr/bin/env python3
"""把群聊存档与频道导出的图片合并成统一图片池（视觉路线的训练/评测入口）。

两个来源的路径基准不同，这里统一解析成绝对路径并补算 dhash，
跨源再去一次重（同一张广告图被转发到不同 chat 的情况确实存在）。
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from typing import Iterable

SOURCES = {
    "group": ("data/processed/archive/image_clf.jsonl", "data/raw/archive/archive"),
    "channel": ("data/processed/channel/images.jsonl", "."),
}


def dhash(path: str, size: int = 8) -> int | None:
    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 load(manifest: str, base: str, source: str, radius: int) -> Iterable[dict]:
    if not os.path.exists(manifest):
        print(f"[skip] 缺 {manifest}", file=sys.stderr)
        return
    with open(manifest, encoding="utf-8") as fh:
        for line in fh:
            r = json.loads(line)
            abs_path = os.path.join(base, r["path"]) if base != "." else os.path.join(os.getcwd(), r["path"])
            if not os.path.exists(abs_path):
                print(f"[warn] 文件缺失 {abs_path}", file=sys.stderr)
                continue
            ph = r.get("dhash")
            ph = int(ph, 16) if isinstance(ph, str) else dhash(abs_path)
            if ph is None:
                continue
            yield {
                "source": source,
                "abs_path": os.path.abspath(abs_path),
                "chat_id": r.get("chat_id"),
                "message_id": r.get("message_id"),
                "date": r.get("date"),
                "label": r["label"],
                "ad_kind": r.get("ad_kind"),
                "label_source": r.get("label_source", "judgment"),
                "width": r.get("width"),
                "height": r.get("height"),
                "sha256": r.get("sha256"),
                "dhash": f"{ph:016x}",
                "_ph": ph,
                "text": (r.get("text") or "")[:300],
            }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/processed/image_pool.jsonl")
    ap.add_argument("--phash-radius", type=int, default=4)
    args = ap.parse_args()

    seen_hash: set[str] = set()
    seen_ph: list[tuple[int, dict]] = []
    pool, dropped = [], {"exact_dup": 0, "near_dup": 0}

    for source, (manifest, base) in SOURCES.items():
        n = 0
        for r in load(manifest, base, source, args.phash_radius):
            n += 1
            if r["sha256"] and r["sha256"] in seen_hash:
                dropped["exact_dup"] += 1
                continue
            if any(((r["_ph"] ^ p).bit_count()) <= args.phash_radius for p, _ in seen_ph):
                dropped["near_dup"] += 1
                continue
            if r["sha256"]:
                seen_hash.add(r["sha256"])
            seen_ph.append((r["_ph"], r))
            r.pop("_ph")
            pool.append(r)
        print(f"[{source}] 读入 {n}")

    pool.sort(key=lambda r: (r["label"] != "ad", r["date"] or 0))
    with open(args.out, "w", encoding="utf-8") as fh:
        for r in pool:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")

    ads = [r for r in pool if r["label"] == "ad"]
    by_source = {}
    for r in ads:
        by_source[r["source"]] = by_source.get(r["source"], 0) + 1
    print(json.dumps({
        "pool": len(pool),
        "dropped": dropped,
        "label_dist": {"ad": len(ads), "clean": len(pool) - len(ads)},
        "ad_by_source": by_source,
        "out": os.path.abspath(args.out),
    }, ensure_ascii=False, indent=2))


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