#!/usr/bin/env python3
"""把引用到的图片复制进 data/processed，让清洗产物自包含。

清洗产物（image_clf.jsonl 等）本身只有路径清单，图片还在原始导出目录里。
删原始数据之前必须先搬，否则清单里的 800 多条路径全部变成死链。

搬完后所有清单里的路径都改写成 processed 内的相对路径，原始目录可安全删除。
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys

POOL_DIR = "data/processed/image_pool"          # 图片落地目录
POOL_MANIFEST = "data/processed/image_pool.jsonl"  # 索引（merge_image_pool.py 的产物）


def sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def dest_rel(source: str, abs_path: str, raw_archive: str, channel_src: str) -> str:
    """把原始绝对路径映射成 processed 内的相对路径，按来源隔离命名空间。"""
    ap = os.path.abspath(abs_path)
    if source == "channel":
        rel = os.path.relpath(ap, os.path.abspath(channel_src))
        return os.path.join("images", "channel", rel.replace("\\", "/")).replace("\\", "/")
    rel = os.path.relpath(ap, os.path.abspath(raw_archive))
    return os.path.join("images", "group", rel.replace("\\", "/")).replace("\\", "/")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw-archive", default="data/raw/archive/archive")
    ap.add_argument("--channel-src", default="存档数据-1")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    with open(POOL_MANIFEST, encoding="utf-8") as fh:
        pool = [json.loads(line) for line in fh]

    # 1) 复制 + 校验
    plan: list[tuple[dict, str, str]] = []
    copied = skipped = mismatch = 0
    for r in pool:
        src = r["abs_path"]
        if not os.path.exists(src):
            print(f"[miss] {src}", file=sys.stderr)
            continue
        rel = dest_rel(r["source"], src, args.raw_archive, args.channel_src)
        dst = os.path.join(POOL_DIR, rel)
        plan.append((r, src, rel))
        if os.path.exists(dst):
            skipped += 1
        else:
            if not args.dry_run:
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                shutil.copy2(src, dst)
            copied += 1
        if not args.dry_run and r.get("sha256") and sha256(dst) != r["sha256"]:
            mismatch += 1
            print(f"[hash-mismatch] {rel}", file=sys.stderr)

    print(f"[copy] 新增 {copied}，已存在 {skipped}，哈希不符 {mismatch}，计划 {len(plan)}")
    if args.dry_run:
        return
    if mismatch:
        print("[abort] 存在哈希不符，放弃改写清单", file=sys.stderr)
        sys.exit(1)

    # 2) 改写 image_pool.jsonl（路径改为 processed 内相对路径）
    for r, _src, rel in plan:
        r["path"] = rel
        r["abs_path"] = os.path.join(os.getcwd(), POOL_DIR, rel).replace("\\", "/")
    with open(POOL_MANIFEST, "w", encoding="utf-8") as fh:
        for r, _s, _rel in plan:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")

    # 3) 改写两个来源清单（archive 用 processed 内绝对路径，channel 用 processed 内相对路径）
    by_key = {(r["source"], r["chat_id"], r["message_id"], os.path.basename(r["abs_path"])): r
              for r, _s, _rel in plan}

    def patch(path: str, key_source: str, fields=("path",)) -> int:
        if not os.path.exists(path):
            return 0
        rows = [json.loads(line) for line in open(path, encoding="utf-8")]
        n = 0
        for row in rows:
            for f in fields:
                old = row.get(f)
                if not isinstance(old, str):
                    continue
                k = (key_source, row.get("chat_id"), row.get("message_id"), os.path.basename(old))
                hit = by_key.get(k)
                if hit:
                    row[f] = os.path.join(os.getcwd(), POOL_DIR, hit["path"]).replace("\\", "/")
                    n += 1
        with open(path, "w", encoding="utf-8") as fh:
            for row in rows:
                fh.write(json.dumps(row, ensure_ascii=False) + "\n")
        return n

    n1 = patch("data/processed/archive/image_clf.jsonl", "group")
    n2 = patch("data/processed/channel/images.jsonl", "channel")
    # channel messages 的 images 是列表
    msgs = [json.loads(line) for line in open("data/processed/channel/messages.jsonl", encoding="utf-8")]
    n3 = 0
    for m in msgs:
        new = []
        for old in m.get("images") or []:
            k = ("channel", m.get("chat_id"), m.get("message_id"), os.path.basename(old))
            hit = by_key.get(k)
            new.append(os.path.join(os.getcwd(), POOL_DIR, hit["path"]).replace("\\", "/") if hit else old)
            n3 += bool(hit)
        m["images"] = new
    with open("data/processed/channel/messages.jsonl", "w", encoding="utf-8") as fh:
        for m in msgs:
            fh.write(json.dumps(m, ensure_ascii=False) + "\n")

    print(f"[patch] image_clf {n1} | channel/images {n2} | channel/messages {n3}")
    print(f"[out] {os.path.abspath(POOL_DIR)}")


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