#!/usr/bin/env python3
"""用 Unity 源贴图 MD5 匹配恢复 __stage__，并完成 sanxing 命名统一。"""
from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
TEX = ROOT / "assets" / "resources" / "textures"
UNITY_TEX = Path("/Users/liuyufei/tfrh/竞赛/scratch-unity-base/Assets/Texture")

# 复用 unify 脚本的 RENAMES
import importlib.util
_spec = importlib.util.spec_from_file_location(
    "unify", ROOT / "tools" / "unify-theme-texture-names.py"
)
_unify = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_unify)
RENAMES = _unify.RENAMES
_norm_rel = _unify._norm_rel
build_rename_lookup = _unify.build_rename_lookup
move_pair = _unify.move_pair
apply_renames = _unify.apply_renames
cleanup_empty_dirs = _unify.cleanup_empty_dirs
move_legacy = _unify.move_legacy
SANXING_KEEP = _unify.SANXING_KEEP
THEME_SHIP = _unify.THEME_SHIP
build_path_replacements = _unify.build_path_replacements
patch_text_files = _unify.patch_text_files
patch_themes_database = _unify.patch_themes_database
patch_palettes = _unify.patch_palettes

UNITY_SOURCES: dict[str, list[Path]] = {
    "silu": [UNITY_TEX / "silu"],
    "numMan": [UNITY_TEX / "numMan"],
    "redArmy": [UNITY_TEX / "redarmy"],
    "chinese": [UNITY_TEX / "Panda", UNITY_TEX / "Chinese"],
    "snow": [TEX / "snow"],  # snow 仅 Cocos，从现有目录取
    "sanxing": [TEX / "sanxing"],
}


def md5_file(path: Path) -> str:
    h = hashlib.md5()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def build_unity_hash_index(theme: str) -> dict[str, str]:
    """md5 -> theme 内相对路径（不含 .png）"""
    index: dict[str, str] = {}
    for src_root in UNITY_SOURCES.get(theme, []):
        if not src_root.exists():
            continue
        for png in src_root.rglob("*.png"):
            rel = png.relative_to(src_root).as_posix().replace(".png", "")
            # chinese: Panda/待机/机器人 -> skin/待机/机器人 便于 lookup
            if theme == "chinese" and rel.startswith("Panda/"):
                rel = "skin/" + rel[len("Panda/"):]
            digest = md5_file(png)
            index[digest] = rel
    return index


def recover_stages_by_hash() -> int:
    lookup = build_rename_lookup()
    recovered = 0
    for stage in sorted(TEX.rglob("__stage_*.png")):
        theme_key = stage.relative_to(TEX).parts[0]
        digest = md5_file(stage)
        rel = build_unity_hash_index(theme_key).get(digest)
        if not rel:
            continue
        new_rel = lookup.get((theme_key, Path(rel).name)) or lookup.get((theme_key, rel))
        if not new_rel:
            continue
        dst = TEX / theme_key / f"{new_rel}.png"
        if move_pair(stage, dst, skip_if_dst_exists=True):
            recovered += 1
            print(f"  hash-recovered {theme_key}: {rel} -> {new_rel}.png")
    return recovered


def copy_from_unity_if_missing(theme: str) -> int:
    """目标路径缺失时从 Unity 复制并重命名"""
    lookup = build_rename_lookup()
    count = 0
    for src_root in UNITY_SOURCES.get(theme, []):
        if not src_root.exists() or src_root == TEX / theme:
            continue
        for png in src_root.rglob("*.png"):
            rel = png.relative_to(src_root).as_posix().replace(".png", "")
            if theme == "chinese" and rel.startswith("Panda/"):
                rel = "skin/" + rel[len("Panda/"):]
            key_name = Path(rel).name
            new_rel = lookup.get((theme, key_name)) or lookup.get((theme, rel))
            if not new_rel:
                continue
            dst = TEX / theme / f"{new_rel}.png"
            if dst.exists():
                continue
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(png, dst)
            count += 1
            print(f"  copied {theme}: {rel} -> {new_rel}.png")
    return count


def remove_orphan_stages() -> None:
    legacy_root = TEX / "_orphan_stages"
    for stage in TEX.rglob("__stage_*.png"):
        rel = stage.relative_to(TEX)
        dest = legacy_root / rel
        dest.parent.mkdir(parents=True, exist_ok=True)
        move_pair(stage, dest)


def main() -> None:
    print("=== MD5 恢复 __stage__ ===")
    n_hash = recover_stages_by_hash()
    print(f"  恢复 {n_hash} 个")

    print("=== 从 Unity 补全缺失文件 ===")
    n_copy = 0
    for theme in ["silu", "numMan", "redArmy", "chinese"]:
        n_copy += copy_from_unity_if_missing(theme)
    print(f"  复制 {n_copy} 个")

    print("=== 常规重命名 ===")
    n = apply_renames()
    print(f"  重命名 {n} 个")

    print("=== 清理 ===")
    remove_orphan_stages()
    for theme in ["silu", "snow", "numMan", "redArmy", "chinese"]:
        cleanup_empty_dirs(TEX / theme)
        keep = set(SANXING_KEEP) | set(THEME_SHIP.get(theme, ()))
        move_legacy(theme, keep)

    print("=== 更新引用 ===")
    patch_text_files(build_path_replacements())
    patch_themes_database()
    patch_palettes()
    print("完成。")


if __name__ == "__main__":
    main()