#!/usr/bin/env python3
"""Restore game assets into pack/restored.

This script rebuilds three major sources into logical paths:
1. Local bundles under assets/assets/*.
2. Incremental bundles under files/GameAssets/assets/*.
3. Remote cache bundles under files/gamecaches/* using cacheList.json.

It also keeps special handling for AnimationSpine and gamecaches/Spine, including:
- beeplay decode
- gzip/ccz inflate
- ASTC -> PNG conversion
- repair of previously mis-decoded PNG headers

Usage:
    python RestoreGameAssets.py
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from collections import defaultdict
import gzip
import io
import json
import os
import re
import shutil
import subprocess
import tempfile
import zlib
import struct


SCRIPT_DIR = Path(__file__).resolve().parent


def env_path(name: str, default: Path) -> Path:
    value = os.environ.get(name, "").strip()
    if not value:
        return default.resolve()

    path = Path(value)
    if not path.is_absolute():
        path = SCRIPT_DIR / path
    return path.resolve()


def env_text(name: str, default: str) -> str:
    value = os.environ.get(name, "").strip()
    return value or default


def env_int(name: str, default: int) -> int:
    value = os.environ.get(name, "").strip()
    if not value:
        return default
    try:
        return int(value, 0)
    except Exception:
        return default


def detect_astcenc_path() -> Path:
    configured = os.environ.get("RESTORE_ASTCENC", "").strip()
    if configured:
        configured_path = Path(configured)
        if not configured_path.is_absolute():
            configured_path = SCRIPT_DIR / configured_path
        return configured_path.resolve()

    candidates = [
        SCRIPT_DIR / "bin" / "astcenc-avx2.exe",
        env_path("RESTORE_ROOT", SCRIPT_DIR) / "bin" / "astcenc-avx2.exe",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return candidates[0]


ROOT = env_path("RESTORE_ROOT", SCRIPT_DIR)
OUTPUT_BASE = env_path("RESTORE_OUTPUT_DIR", ROOT / "pack" / "restored")
BACKUP_BASE = OUTPUT_BASE / "_partial_backups"
APK_ASSETS_ROOT = env_path("RESTORE_ASSETS_ROOT", ROOT / "assets" / "assets")
GAMEASSETS_ROOT = env_path("RESTORE_GAMEASSETS_ROOT", ROOT / "files" / "GameAssets" / "assets")
GAMECACHE_ROOT = env_path("RESTORE_GAMECACHE_ROOT", ROOT / "files" / "gamecaches")
ANIMATION_SPINE_BUNDLE = env_text("RESTORE_ANIMATION_SPINE_BUNDLE", "AnimationSpine")
SPINE_CACHE_BUNDLE = env_text("RESTORE_SPINE_BUNDLE", "Spine")
GAMECACHE_CACHELIST = GAMECACHE_ROOT / "cacheList.json"
DEFAULT_SPINE_CONFIG = GAMECACHE_ROOT / SPINE_CACHE_BUNDLE / "177385225976168.json"
ASTCENC = detect_astcenc_path()
ASTC_MODE = "s"

PREBUILT_OUTPUTS = [
    (APK_ASSETS_ROOT / ANIMATION_SPINE_BUNDLE / "output", OUTPUT_BASE / ANIMATION_SPINE_BUNDLE),
]

EXACT_NATIVE_BUNDLES = [
    (
        f"{ANIMATION_SPINE_BUNDLE}/GameAssets",
        GAMEASSETS_ROOT / ANIMATION_SPINE_BUNDLE / "config.json",
        GAMEASSETS_ROOT / ANIMATION_SPINE_BUNDLE / "native",
        OUTPUT_BASE / ANIMATION_SPINE_BUNDLE,
    ),
]

LOCAL_BUNDLE_ROOTS = [
    APK_ASSETS_ROOT,
    GAMEASSETS_ROOT,
]

SPECIAL_LOCAL_BUNDLES = {ANIMATION_SPINE_BUNDLE}
SPECIAL_GAMECACHE_BUNDLES = {SPINE_CACHE_BUNDLE}
ONLY_BUNDLES = {
    item.strip()
    for item in os.environ.get("RESTORE_ONLY_BUNDLES", "").split(",")
    if item.strip()
}

UUID_RE = re.compile(
    r"/([0-9a-f-]{8,36})\.[0-9a-f]+\.(json|atlas|bin|astc|png|jpg)$",
    re.IGNORECASE,
)

SIGNATURE = env_text("RESTORE_BEEPLAY_SIGNATURE", "beeplay").encode("utf-8")
XOR_KEY_DEFAULT = env_int("RESTORE_XOR_KEY_DEFAULT", 0x17) & 0xFF
XOR_KEY_IMAGE = env_int("RESTORE_XOR_KEY_IMAGE", 0x44) & 0xFF
XOR_START = 16
XOR_END = 24
BASE64_KEYS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
HEX_CHARS = "0123456789abcdef"
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
JPEG_SIGNATURE = b"\xff\xd8\xff"
WEBP_SIGNATURE = b"RIFF"
WRONG_XOR_DELTA = XOR_KEY_DEFAULT ^ XOR_KEY_IMAGE

JSON_TYPE_PRIORITY = {
    "sp.SkeletonData": 0,
    "cc.SpriteFrame": 1,
    "cc.SpriteAtlas": 2,
    "cc.ParticleAsset": 3,
    "cc.Prefab": 4,
    "cc.AnimationClip": 5,
    "cc.Material": 6,
    "cc.EffectAsset": 7,
    "cc.TextAsset": 8,
    "cc.Asset": 9,
    "cc.Texture2D": 10,
}


@dataclass
class SpineAsset:
    logical_path: str
    atlas: Path | None = None
    skeleton: Path | None = None
    images: list[Path] = field(default_factory=list)
    skeleton_json: Path | None = None
    atlas_json: Path | None = None


def load_json(path: Path):
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None


def safe_write_bytes(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(data)


def display_path(path: Path) -> str:
    try:
        return str(path.relative_to(ROOT))
    except Exception:
        return str(path)


def is_probable_bundle_config(path: Path) -> bool:
    obj = load_json(path)
    return isinstance(obj, dict) and isinstance(obj.get("paths"), dict) and isinstance(obj.get("uuids"), list)


def discover_spine_config() -> Path | None:
    configured = os.environ.get("RESTORE_SPINE_CONFIG", "").strip()
    if configured:
        configured_path = Path(configured)
        if not configured_path.is_absolute():
            configured_path = ROOT / configured_path
        if configured_path.exists() and is_probable_bundle_config(configured_path):
            return configured_path

    if DEFAULT_SPINE_CONFIG.exists() and is_probable_bundle_config(DEFAULT_SPINE_CONFIG):
        return DEFAULT_SPINE_CONFIG

    spine_root = GAMECACHE_ROOT / SPINE_CACHE_BUNDLE
    if not spine_root.exists():
        return None

    candidates = sorted(spine_root.glob("*.json"), key=lambda item: item.name.lower())
    for candidate in candidates:
        if is_probable_bundle_config(candidate):
            return candidate

    return None


def png_ihdr_crc_valid(data: bytes) -> bool:
    if len(data) < 33 or not data.startswith(PNG_SIGNATURE):
        return False
    if data[12:16] != b"IHDR":
        return False

    expected = struct.unpack(">I", data[29:33])[0]
    actual = zlib.crc32(data[12:29]) & 0xFFFFFFFF
    return actual == expected


def maybe_repair_wrong_xor_png(data: bytes) -> bytes:
    if not data.startswith(PNG_SIGNATURE) or png_ihdr_crc_valid(data):
        return data
    if len(data) < 24:
        return data

    candidate_ranges = [
        range(16, 24),
        list(range(9, 16)) + list(range(17, 24)),
    ]

    for index_range in candidate_ranges:
        repaired = bytearray(data)
        for index in index_range:
            repaired[index] ^= WRONG_XOR_DELTA
        repaired_bytes = bytes(repaired)
        if png_ihdr_crc_valid(repaired_bytes):
            return repaired_bytes

    return data


def sync_tree(src: Path, dst: Path) -> int:
    if not src.exists():
        return 0

    copied = 0
    for file_path in src.rglob("*"):
        if not file_path.is_file():
            continue
        dest_path = dst / file_path.relative_to(src)
        data = file_path.read_bytes()
        if file_path.suffix.lower() == ".png":
            data = maybe_repair_wrong_xor_png(data)
        safe_write_bytes(dest_path, data)
        copied += 1

    return copied


def is_gzip(data: bytes) -> bool:
    return data.startswith(b"\x1f\x8b")


def is_ccz(data: bytes) -> bool:
    return data.startswith(b"CCZ!")


def inflate_gzip(data: bytes) -> bytes:
    with gzip.GzipFile(fileobj=io.BytesIO(data)) as file_obj:
        return file_obj.read()


def inflate_ccz(data: bytes) -> bytes:
    return zlib.decompress(data[16:])


def decrypt_beeplay(data: bytes, key: int = XOR_KEY_DEFAULT) -> bytes:
    buf = bytearray(data)
    for i in range(XOR_START, min(XOR_END, len(buf))):
        buf[i] ^= key
    return bytes(buf)


def pick_beeplay_key(data: bytes) -> int:
    if data.startswith(PNG_SIGNATURE):
        return XOR_KEY_IMAGE
    if data.startswith(JPEG_SIGNATURE):
        return XOR_KEY_IMAGE
    if data.startswith(WEBP_SIGNATURE) and data[8:12] == b"WEBP":
        return XOR_KEY_IMAGE
    return XOR_KEY_DEFAULT


def decode_blob(data: bytes) -> bytes:
    if is_ccz(data):
        data = inflate_ccz(data)
    elif is_gzip(data):
        data = inflate_gzip(data)

    if data.startswith(SIGNATURE):
        payload = data[len(SIGNATURE) :]
        data = decrypt_beeplay(payload, pick_beeplay_key(payload))

    return data


def decompress_cocos_uuid(uuid_value: str) -> str:
    if len(uuid_value) == 36 and uuid_value.count("-") == 4:
        return uuid_value.lower()

    head = uuid_value[:2]
    rest = uuid_value[2:]
    hex_chars = [head[0], head[1]]

    for i in range(0, len(rest), 2):
        if i + 1 >= len(rest):
            break
        left = BASE64_KEYS.index(rest[i])
        right = BASE64_KEYS.index(rest[i + 1])
        a = left >> 2
        b = ((left & 0x3) << 2) | (right >> 4)
        c = right & 0xF
        hex_chars.extend([HEX_CHARS[a], HEX_CHARS[b], HEX_CHARS[c]])
        if len(hex_chars) >= 32:
            break

    full = "".join(hex_chars[:32])
    return f"{full[:8]}-{full[8:12]}-{full[12:16]}-{full[16:20]}-{full[20:]}".lower()


def build_uuid_to_logical_path(config_path: Path) -> dict[str, str]:
    config = load_json(config_path)
    if not isinstance(config, dict):
        return {}

    paths = config.get("paths", {}) or {}
    uuids = config.get("uuids", []) or []
    mapping: dict[str, str] = {}

    for key, info in paths.items():
        if not isinstance(info, list) or not info:
            continue
        if not isinstance(info[0], str):
            continue

        index = int(key)
        if index >= len(uuids):
            continue

        uuid36 = decompress_cocos_uuid(uuids[index])
        mapping[uuid36] = info[0]

    return mapping


def cache_sort_key(path: Path):
    stem = path.stem
    return int(stem) if stem.isdigit() else stem


def parse_import_json_kind(path: Path) -> str:
    obj = load_json(path)
    if not isinstance(obj, list) or len(obj) <= 5 or not isinstance(obj[3], list):
        return "other"

    descriptors = [item[0] for item in obj[3] if isinstance(item, list) and item]
    if "sp.SkeletonData" in descriptors:
        return "skeleton"
    if "cc.Asset" in descriptors:
        return "atlas"
    return "other"


def extract_skeleton_atlas_text(path: Path) -> str:
    obj = load_json(path)
    if not isinstance(obj, list) or len(obj) <= 5:
        return ""

    for item in obj[5]:
        if (
            isinstance(item, list)
            and len(item) >= 4
            and isinstance(item[2], str)
            and item[2] in (".bin", ".skel")
            and isinstance(item[3], str)
        ):
            return item[3].lstrip("\r\n")

    return ""


def read_text_file(path: Path) -> str:
    try:
        return decode_blob(path.read_bytes()).decode("utf-8", errors="ignore")
    except Exception:
        return ""


def parse_atlas_texture_names(text: str) -> list[str]:
    names: list[str] = []
    expect_texture = True

    for line in text.splitlines():
        line = line.strip()
        if not line:
            expect_texture = True
            continue
        if line.startswith("#"):
            continue
        if expect_texture:
            names.append(line)
            expect_texture = False

    return names


def convert_astc_bytes_to_png(astc_bytes: bytes, png_path: Path) -> bool:
    if not ASTCENC.exists():
        return False

    png_path.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory(prefix="restore_spine_", dir=ROOT) as tmp_dir:
        tmp_dir_path = Path(tmp_dir)
        temp_astc = tmp_dir_path / "temp.astc"
        temp_png = tmp_dir_path / "temp.png"
        temp_astc.write_bytes(astc_bytes)

        try:
            subprocess.run(
                [str(ASTCENC), f"-d{ASTC_MODE}", str(temp_astc), str(temp_png)],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            shutil.copy2(temp_png, png_path)
            return True
        except Exception:
            return False


def restore_image(src_path: Path, dst_path: Path) -> bool:
    data = decode_blob(src_path.read_bytes())
    ext = src_path.suffix.lower()

    if ext == ".astc":
        if convert_astc_bytes_to_png(data, dst_path):
            return True
        fallback = dst_path.with_suffix(".astc")
        safe_write_bytes(fallback, data)
        return False

    safe_write_bytes(dst_path, data)
    return True


def collect_gamecache_spine_assets() -> tuple[dict[str, SpineAsset], int]:
    cache_list = load_json(GAMECACHE_CACHELIST)
    spine_config = discover_spine_config()
    if spine_config is None:
        print(f"[gamecaches/{SPINE_CACHE_BUNDLE}] WARN: no usable Spine config json was found")
        return {}, 0

    uuid_to_path = build_uuid_to_logical_path(spine_config)
    if not isinstance(cache_list, dict):
        return {}, 0

    entries = cache_list.get("files", {}) or {}
    assets: dict[str, SpineAsset] = {}
    missing_mappings = 0

    for remote_path, meta in entries.items():
        if not isinstance(remote_path, str) or not remote_path.startswith(f"remote/{SPINE_CACHE_BUNDLE}/"):
            continue
        if not isinstance(meta, dict):
            continue

        match = UUID_RE.search(remote_path.replace("\\", "/"))
        if not match:
            continue

        uuid36 = match.group(1).lower()
        logical_path = uuid_to_path.get(uuid36)
        if not logical_path:
            missing_mappings += 1
            continue

        relative_cached = meta.get("url")
        if not isinstance(relative_cached, str):
            continue

        cached_path = GAMECACHE_ROOT / relative_cached
        if not cached_path.exists():
            continue

        asset = assets.setdefault(logical_path, SpineAsset(logical_path=logical_path))

        if "/native/" in remote_path:
            suffix = cached_path.suffix.lower()
            if suffix == ".atlas":
                asset.atlas = cached_path
            elif suffix in (".bin", ".skel"):
                asset.skeleton = cached_path
            elif suffix in (".astc", ".png", ".jpg"):
                asset.images.append(cached_path)
        elif "/import/" in remote_path and cached_path.suffix.lower() == ".json":
            kind = parse_import_json_kind(cached_path)
            if kind == "skeleton" and asset.skeleton_json is None:
                asset.skeleton_json = cached_path
            elif kind == "atlas" and asset.atlas_json is None:
                asset.atlas_json = cached_path

    for asset in assets.values():
        asset.images.sort(key=cache_sort_key)

    return assets, missing_mappings


def collect_native_bundle_assets(config_path: Path, native_root: Path) -> dict[str, SpineAsset]:
    uuid_to_path = build_uuid_to_logical_path(config_path)
    assets: dict[str, SpineAsset] = {}

    for file_path in native_root.rglob("*"):
        if not file_path.is_file():
            continue

        suffix = file_path.suffix.lower()
        if suffix not in {".atlas", ".bin", ".skel", ".astc", ".png", ".jpg"}:
            continue

        uuid36 = file_path.stem.lower()
        logical_path = uuid_to_path.get(uuid36)
        if not logical_path:
            continue

        asset = assets.setdefault(logical_path, SpineAsset(logical_path=logical_path))
        if suffix == ".atlas":
            asset.atlas = file_path
        elif suffix in {".bin", ".skel"}:
            asset.skeleton = file_path
        elif suffix in {".astc", ".png", ".jpg"}:
            asset.images.append(file_path)

    for asset in assets.values():
        asset.images.sort(key=cache_sort_key)

    return assets


def backup_existing_dir(target_dir: Path) -> Path | None:
    if not target_dir.exists():
        return None

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    try:
        relative_target = target_dir.relative_to(OUTPUT_BASE)
    except Exception:
        relative_target = Path(target_dir.name)

    backup_parent = BACKUP_BASE / relative_target.parent
    backup_parent.mkdir(parents=True, exist_ok=True)
    backup_dir = backup_parent / f"{target_dir.name}_partial_backup_{timestamp}"
    shutil.move(str(target_dir), str(backup_dir))
    return backup_dir


def migrate_legacy_backups() -> None:
    if not OUTPUT_BASE.exists():
        return

    legacy_dirs: list[Path] = []
    for candidate in OUTPUT_BASE.rglob("*_partial_backup_*"):
        if not candidate.is_dir():
            continue
        if BACKUP_BASE in candidate.parents:
            continue
        legacy_dirs.append(candidate)

    for legacy_dir in sorted(legacy_dirs, key=lambda item: len(item.parts)):
        try:
            relative_legacy = legacy_dir.relative_to(OUTPUT_BASE)
        except Exception:
            continue

        backup_parent = BACKUP_BASE / relative_legacy.parent
        backup_parent.mkdir(parents=True, exist_ok=True)
        target_dir = backup_parent / legacy_dir.name
        if target_dir.exists():
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            target_dir = backup_parent / f"{legacy_dir.name}_{timestamp}"
        shutil.move(str(legacy_dir), str(target_dir))
        print(f"[backup] moved legacy snapshot to: {display_path(target_dir)}")


def cleanup_stale_stage_dirs() -> None:
    if not OUTPUT_BASE.exists():
        return

    stale_dirs: list[Path] = []
    for candidate in OUTPUT_BASE.rglob("*"):
        if not candidate.is_dir():
            continue
        if candidate.name.startswith("_") and candidate.name.endswith("_stage"):
            stale_dirs.append(candidate)

    for stale_dir in sorted(stale_dirs, key=lambda item: len(item.parts), reverse=True):
        shutil.rmtree(stale_dir, ignore_errors=True)
        print(f"[cleanup] removed stale stage dir: {display_path(stale_dir)}")


def restore_spine_asset(asset: SpineAsset, output_root: Path) -> tuple[bool, bool, bool]:
    logical_path = Path(asset.logical_path)
    dest_dir = output_root / logical_path.parent
    base_name = logical_path.name
    dest_dir.mkdir(parents=True, exist_ok=True)

    restored_atlas = False
    restored_skeleton = False
    restored_image = False

    atlas_text = read_text_file(asset.atlas) if asset.atlas else ""
    if not atlas_text and asset.skeleton_json:
        atlas_text = extract_skeleton_atlas_text(asset.skeleton_json)

    if atlas_text:
        safe_write_bytes((dest_dir / f"{base_name}.atlas"), atlas_text.encode("utf-8"))
        restored_atlas = True

    if asset.skeleton:
        safe_write_bytes(dest_dir / f"{base_name}.skel", decode_blob(asset.skeleton.read_bytes()))
        restored_skeleton = True

    if asset.skeleton_json:
        safe_write_bytes(dest_dir / f"{base_name}.json", decode_blob(asset.skeleton_json.read_bytes()))

    image_targets = parse_atlas_texture_names(atlas_text)
    if asset.images:
        if not image_targets:
            image_targets = [f"{base_name}.png"]
        if len(image_targets) < len(asset.images):
            for index in range(len(image_targets), len(asset.images)):
                image_targets.append(f"{base_name}_{index + 1}.png")

        # Prefer already-decoded png/jpg over astc if the same logical asset was cached twice.
        images = sorted(asset.images, key=lambda p: (p.suffix.lower() == ".astc", cache_sort_key(p)))
        for src_path, target_name in zip(images, image_targets):
            target_path = dest_dir / target_name
            target_path.parent.mkdir(parents=True, exist_ok=True)
            if restore_image(src_path, target_path):
                restored_image = True

    return restored_atlas, restored_skeleton, restored_image


def restore_gamecache_spine() -> None:
    assets, missing_mappings = collect_gamecache_spine_assets()
    label = f"gamecaches/{SPINE_CACHE_BUNDLE}"
    if not assets:
        print(f"[{label}] skipped: no logical assets were collected")
        return

    stage_dir = OUTPUT_BASE / "gamecaches" / f"_{SPINE_CACHE_BUNDLE}_stage"
    target_dir = OUTPUT_BASE / "gamecaches" / SPINE_CACHE_BUNDLE

    if stage_dir.exists():
        shutil.rmtree(stage_dir)
    stage_dir.mkdir(parents=True, exist_ok=True)

    atlas_count = 0
    skeleton_count = 0
    image_count = 0

    for asset in sorted(assets.values(), key=lambda item: item.logical_path.lower()):
        restored_atlas, restored_skeleton, restored_image = restore_spine_asset(asset, stage_dir)
        atlas_count += int(restored_atlas)
        skeleton_count += int(restored_skeleton)
        image_count += int(restored_image)

    backup_dir = backup_existing_dir(target_dir)
    stage_dir.rename(target_dir)

    print(f"[{label}] rebuilt {len(assets)} logical assets")
    print(f"[{label}] atlas={atlas_count} skel={skeleton_count} image={image_count}")
    if missing_mappings:
        print(f"[{label}] WARN: {missing_mappings} cache entries had no config mapping")
    if backup_dir:
        print(f"[{label}] previous output moved to: {display_path(backup_dir)}")


def restore_native_bundle(label: str, config_path: Path, native_root: Path, output_root: Path) -> int:
    if not config_path.exists() or not native_root.exists():
        return 0

    assets = collect_native_bundle_assets(config_path, native_root)
    atlas_count = 0
    skeleton_count = 0
    image_count = 0

    for asset in sorted(assets.values(), key=lambda item: item.logical_path.lower()):
        restored_atlas, restored_skeleton, restored_image = restore_spine_asset(asset, output_root)
        atlas_count += int(restored_atlas)
        skeleton_count += int(restored_skeleton)
        image_count += int(restored_image)

    print(f"[{label}] restored {len(assets)} logical assets")
    print(f"[{label}] atlas={atlas_count} skel={skeleton_count} image={image_count}")
    return len(assets)


def type_priority(type_name: str) -> int:
    return JSON_TYPE_PRIORITY.get(type_name, 100)


def candidate_names(value: str) -> list[str]:
    value = value.strip()
    if not value:
        return []

    names: list[str] = []
    seen: set[str] = set()

    def add(name: str) -> None:
        lowered = name.lower()
        if lowered not in seen:
            seen.add(lowered)
            names.append(lowered)

    add(value)
    if len(value) >= 20 and value.count("-") != 4:
        try:
            add(decompress_cocos_uuid(value))
        except Exception:
            pass

    return names


def build_file_index(root: Path) -> dict[str, list[Path]]:
    index: dict[str, list[Path]] = defaultdict(list)
    if not root.exists():
        return index

    for file_path in root.rglob("*"):
        if not file_path.is_file():
            continue
        index[file_path.stem.lower()].append(file_path)

    return index


def choose_indexed_path(index: dict[str, list[Path]], names: list[str], native: bool = False) -> Path | None:
    matches: list[Path] = []
    seen: set[str] = set()

    for name in names:
        for file_path in index.get(name.lower(), []):
            key = str(file_path).lower()
            if key in seen:
                continue
            seen.add(key)
            matches.append(file_path)

    if not matches:
        return None

    if native:
        extension_order = {
            ".png": 0,
            ".jpg": 1,
            ".jpeg": 1,
            ".webp": 2,
            ".astc": 3,
            ".plist": 4,
            ".atlas": 5,
        }
        matches.sort(key=lambda item: (extension_order.get(item.suffix.lower(), 10), item.name.lower()))
    else:
        matches.sort(key=lambda item: item.name.lower())

    return matches[0]


def load_decoded_json_file(path: Path):
    try:
        text = decode_blob(path.read_bytes()).decode("utf-8", errors="ignore")
        return json.loads(text)
    except Exception:
        return None


def dump_json_bytes(obj) -> bytes:
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":")).encode("utf-8")


def extract_texture_setter_names(payload) -> list[str]:
    if not (
        isinstance(payload, list)
        and len(payload) > 2
        and isinstance(payload[1], list)
        and isinstance(payload[2], list)
    ):
        return []

    names: list[str] = []
    seen: set[str] = set()
    dependencies = payload[1]
    properties = payload[2]

    for index, property_name in enumerate(properties):
        if property_name != "_textureSetter":
            continue
        if index >= len(dependencies):
            continue
        dependency = dependencies[index]
        if not isinstance(dependency, str):
            continue
        for candidate in candidate_names(dependency):
            if candidate not in seen:
                seen.add(candidate)
                names.append(candidate)

    return names


def extract_pack_member(pack_obj, member_index: int, member_count: int):
    if isinstance(pack_obj, list) and len(pack_obj) >= 6 and isinstance(pack_obj[5], list):
        members = pack_obj[5]
        if member_index < len(members):
            member = members[member_index]
            if isinstance(member, list):
                return pack_obj[:5] + member
            return pack_obj[:5] + [member]
        return pack_obj

    if isinstance(pack_obj, dict):
        data = pack_obj.get("data")
        if isinstance(data, str):
            segments = data.split("|")
            if len(segments) == member_count and member_index < len(segments):
                extracted = dict(pack_obj)
                extracted["data"] = segments[member_index]
                return extracted
        return pack_obj

    return pack_obj


def copy_bundle_root_files(bundle_root: Path, output_root: Path) -> None:
    for file_path in bundle_root.iterdir():
        if not file_path.is_file():
            continue
        safe_write_bytes(output_root / file_path.name, decode_blob(file_path.read_bytes()))


def copy_named_root_files(root_files: dict[str, Path], output_root: Path) -> None:
    for relative_name, file_path in root_files.items():
        safe_write_bytes(output_root / relative_name, decode_blob(file_path.read_bytes()))


def replace_target_dir(target_dir: Path, stage_dir: Path) -> Path | None:
    if target_dir.exists():
        if any(target_dir.iterdir()):
            backup_dir = backup_existing_dir(target_dir)
        else:
            shutil.rmtree(target_dir)
            backup_dir = None
    else:
        backup_dir = None

    stage_dir.rename(target_dir)
    return backup_dir


def restore_import_payload(payload, logical_path: str, output_root: Path) -> bool:
    if payload is None:
        return False

    output_path = (output_root / logical_path).with_suffix(".json")
    if output_path.exists():
        return False

    if isinstance(payload, bytes):
        data = payload
    elif isinstance(payload, str):
        data = payload.encode("utf-8")
    else:
        data = dump_json_bytes(payload)

    safe_write_bytes(output_path, data)
    return True


def restore_native_file(src_path: Path, logical_path: str, output_root: Path) -> bool:
    base_path = output_root / logical_path
    suffix = src_path.suffix.lower()

    if suffix == ".astc":
        target_path = base_path.with_suffix(".png")
        return restore_image(src_path, target_path)

    target_path = base_path.with_suffix(suffix)

    data = decode_blob(src_path.read_bytes())
    if suffix == ".png":
        data = maybe_repair_wrong_xor_png(data)
    if suffix == ".atlas":
        try:
            data = data.decode("utf-8", errors="ignore").replace(".astc", ".png").encode("utf-8")
        except Exception:
            pass

    safe_write_bytes(target_path, data)
    return True


def restore_indexed_bundle_source(
    label: str,
    config_path: Path,
    import_index: dict[str, list[Path]],
    native_index: dict[str, list[Path]],
    root_files: dict[str, Path],
    output_root: Path,
) -> tuple[int, int]:
    config = load_json(config_path)
    if not isinstance(config, dict):
        config = load_decoded_json_file(config_path)
    if not isinstance(config, dict):
        return 0, 0

    copy_named_root_files(root_files, output_root)

    types = config.get("types", []) or []
    uuids = config.get("uuids", []) or []
    paths = config.get("paths", {}) or {}
    packs = config.get("packs", {}) or {}

    pack_positions: dict[int, tuple[str, int, int]] = {}
    for pack_name, indices in packs.items():
        if not isinstance(indices, list):
            continue
        for member_index, raw_idx in enumerate(indices):
            try:
                asset_index = int(raw_idx)
            except Exception:
                continue
            pack_positions[asset_index] = (str(pack_name), member_index, len(indices))

    logical_type_sets: dict[str, set[str]] = defaultdict(set)
    for raw_idx, info in paths.items():
        if not isinstance(info, list) or len(info) < 2 or not isinstance(info[0], str):
            continue
        try:
            type_name = types[info[1]]
        except Exception:
            type_name = "unknown"
        logical_type_sets[info[0]].add(type_name)

    json_cache: dict[Path, object] = {}
    records: list[tuple[str, int, str, object, Path | None, set[str]]] = []
    atlas_paths_by_dir: dict[str, list[str]] = defaultdict(list)
    spriteframe_textures_by_dir: dict[str, list[str]] = defaultdict(list)

    for raw_idx, info in paths.items():
        if not isinstance(info, list) or len(info) < 2 or not isinstance(info[0], str):
            continue

        try:
            asset_index = int(raw_idx)
        except Exception:
            continue

        logical_path = info[0]
        try:
            type_name = types[info[1]]
        except Exception:
            type_name = "unknown"

        asset_uuid = uuids[asset_index] if asset_index < len(uuids) else ""
        import_payload = None

        pack_meta = pack_positions.get(asset_index)
        if pack_meta:
            pack_name, member_index, member_count = pack_meta
            pack_path = choose_indexed_path(import_index, [pack_name.lower()])
            if pack_path:
                pack_obj = json_cache.get(pack_path)
                if pack_obj is None:
                    pack_obj = load_decoded_json_file(pack_path)
                    json_cache[pack_path] = pack_obj
                if pack_obj is not None:
                    import_payload = extract_pack_member(pack_obj, member_index, member_count)
        else:
            import_path = choose_indexed_path(import_index, candidate_names(asset_uuid))
            if import_path:
                import_payload = json_cache.get(import_path)
                if import_payload is None:
                    import_payload = load_decoded_json_file(import_path)
                    json_cache[import_path] = import_payload

        native_path = None
        dependency_names: list[str] = []
        if isinstance(import_payload, list) and len(import_payload) > 1 and isinstance(import_payload[1], list):
            for item in import_payload[1]:
                if isinstance(item, str):
                    dependency_names.extend(candidate_names(item))

        search_names = candidate_names(asset_uuid)
        if type_name == "cc.SpriteFrame":
            if "cc.Texture2D" in logical_type_sets[logical_path]:
                search_names = dependency_names + search_names
            else:
                search_names = []
        else:
            search_names = search_names + dependency_names

        if search_names:
            native_path = choose_indexed_path(native_index, search_names, native=True)

        logical_dir = str(Path(logical_path).parent).replace("\\", "/")
        if type_name == "cc.SpriteAtlas":
            atlas_paths_by_dir[logical_dir].append(logical_path)
        elif type_name == "cc.SpriteFrame":
            for texture_name in extract_texture_setter_names(import_payload):
                spriteframe_textures_by_dir[logical_dir].append(texture_name)

        records.append(
            (
                logical_path,
                type_priority(type_name),
                type_name,
                import_payload,
                native_path,
                logical_type_sets[logical_path],
            )
        )

    records.sort(key=lambda item: (item[0].lower(), item[1], item[2].lower()))

    json_count = 0
    native_count = 0

    for logical_path, _, type_name, import_payload, native_path, type_set in records:
        if restore_import_payload(import_payload, logical_path, output_root):
            json_count += 1

        if type_name == "cc.SpriteFrame" and "cc.Texture2D" not in type_set:
            continue
        if native_path and restore_native_file(native_path, logical_path, output_root):
            native_count += 1

    for logical_dir, atlas_paths in atlas_paths_by_dir.items():
        if len(atlas_paths) != 1:
            continue

        texture_names = list(dict.fromkeys(spriteframe_textures_by_dir.get(logical_dir, [])))
        if len(texture_names) != 1:
            continue

        atlas_logical_path = atlas_paths[0]
        atlas_output_path = (output_root / atlas_logical_path).with_suffix(".png")
        if atlas_output_path.exists():
            continue

        inferred_native_path = choose_indexed_path(native_index, texture_names, native=True)
        if inferred_native_path and restore_native_file(inferred_native_path, atlas_logical_path, output_root):
            native_count += 1
            print(f"[{label}] inferred atlas texture for {atlas_logical_path} from {inferred_native_path.name}")

    print(f"[{label}] json={json_count} native={native_count}")
    return json_count, native_count


def restore_local_bundle_source(label: str, bundle_root: Path, output_root: Path) -> tuple[int, int]:
    root_files = {
        file_path.name: file_path
        for file_path in bundle_root.iterdir()
        if file_path.is_file()
    }

    return restore_indexed_bundle_source(
        label=label,
        config_path=bundle_root / "config.json",
        import_index=build_file_index(bundle_root / "import"),
        native_index=build_file_index(bundle_root / "native"),
        root_files=root_files,
        output_root=output_root,
    )


def restore_local_bundles() -> None:
    bundle_names: set[str] = set()
    for root_path in LOCAL_BUNDLE_ROOTS:
        if not root_path.exists():
            continue
        for config_path in root_path.glob("*/config.json"):
            bundle_names.add(config_path.parent.name)

    for bundle_name in sorted(bundle_names):
        if bundle_name in SPECIAL_LOCAL_BUNDLES:
            continue
        if ONLY_BUNDLES and bundle_name not in ONLY_BUNDLES:
            continue

        target_dir = OUTPUT_BASE / bundle_name
        stage_dir = OUTPUT_BASE / f"_{bundle_name}_stage"

        if stage_dir.exists():
            shutil.rmtree(stage_dir)
        stage_dir.mkdir(parents=True, exist_ok=True)

        restored_any = False

        for root_path in LOCAL_BUNDLE_ROOTS:
            bundle_root = root_path / bundle_name
            if not bundle_root.exists():
                continue

            restored_any = True
            output_dir = bundle_root / "output"
            if output_dir.exists():
                copied = sync_tree(output_dir, stage_dir)
                if copied:
                    print(f"[sync/{bundle_name}] copied {copied} files from {output_dir} -> {stage_dir}")

            restore_local_bundle_source(f"{bundle_name}/{bundle_root.parent.parent.name}", bundle_root, stage_dir)

        if not restored_any:
            shutil.rmtree(stage_dir)
            continue

        backup_dir = replace_target_dir(target_dir, stage_dir)
        if backup_dir:
            print(f"[{bundle_name}] previous output moved to: {display_path(backup_dir)}")


def normalize_remote_output_name(relative_remote_path: str) -> str:
    relative_remote_path = relative_remote_path.replace("\\", "/")
    remote_path = Path(relative_remote_path)
    parts = remote_path.name.rsplit(".", 2)
    normalized_name = f"{parts[0]}.{parts[2]}" if len(parts) == 3 else remote_path.name
    if str(remote_path.parent) in ("", "."):
        return normalized_name
    return str(remote_path.parent / normalized_name)


def extract_remote_stem(relative_remote_path: str) -> str:
    basename = Path(relative_remote_path.replace("\\", "/")).name
    parts = basename.rsplit(".", 2)
    return parts[0] if len(parts) == 3 else Path(basename).stem


def build_gamecache_bundle_source(
    bundle_name: str, cache_entries: dict[str, dict]
) -> tuple[Path | None, dict[str, list[Path]], dict[str, list[Path]], dict[str, Path]]:
    prefix = f"remote/{bundle_name}/"
    import_index: dict[str, list[Path]] = defaultdict(list)
    native_index: dict[str, list[Path]] = defaultdict(list)
    root_files: dict[str, Path] = {}
    config_path: Path | None = None

    for remote_path, meta in cache_entries.items():
        if not isinstance(remote_path, str) or not remote_path.startswith(prefix):
            continue
        if not isinstance(meta, dict):
            continue

        relative_cached = meta.get("url")
        if not isinstance(relative_cached, str):
            continue

        cached_path = GAMECACHE_ROOT / relative_cached
        if not cached_path.exists():
            continue

        relative_remote = remote_path[len(prefix) :].replace("\\", "/")
        if relative_remote.startswith("import/"):
            import_index[extract_remote_stem(relative_remote).lower()].append(cached_path)
        elif relative_remote.startswith("native/"):
            native_index[extract_remote_stem(relative_remote).lower()].append(cached_path)
        else:
            output_name = normalize_remote_output_name(relative_remote)
            root_files[output_name] = cached_path
            if output_name.lower() == "config.json":
                config_path = cached_path

    return config_path, import_index, native_index, root_files


def restore_gamecache_bundles() -> None:
    cache_list = load_json(GAMECACHE_CACHELIST)
    if not isinstance(cache_list, dict):
        return

    cache_entries = cache_list.get("files", {}) or {}
    bundle_names = sorted(path.name for path in GAMECACHE_ROOT.iterdir() if path.is_dir())

    for bundle_name in bundle_names:
        if bundle_name in SPECIAL_GAMECACHE_BUNDLES:
            continue
        if ONLY_BUNDLES and bundle_name not in ONLY_BUNDLES and f"gamecaches/{bundle_name}" not in ONLY_BUNDLES:
            continue

        config_path, import_index, native_index, root_files = build_gamecache_bundle_source(bundle_name, cache_entries)
        if config_path is None:
            continue

        target_dir = OUTPUT_BASE / "gamecaches" / bundle_name
        stage_dir = OUTPUT_BASE / "gamecaches" / f"_{bundle_name}_stage"

        if stage_dir.exists():
            shutil.rmtree(stage_dir)
        stage_dir.mkdir(parents=True, exist_ok=True)

        restore_indexed_bundle_source(
            label=f"gamecaches/{bundle_name}",
            config_path=config_path,
            import_index=import_index,
            native_index=native_index,
            root_files=root_files,
            output_root=stage_dir,
        )

        backup_dir = replace_target_dir(target_dir, stage_dir)
        if backup_dir:
            print(f"[gamecaches/{bundle_name}] previous output moved to: {display_path(backup_dir)}")


def restore_animation_spine() -> None:
    target_dir = OUTPUT_BASE / ANIMATION_SPINE_BUNDLE
    stage_dir = OUTPUT_BASE / f"_{ANIMATION_SPINE_BUNDLE}_stage"

    if stage_dir.exists():
        shutil.rmtree(stage_dir)
    stage_dir.mkdir(parents=True, exist_ok=True)

    restored_any = False

    for src, dst in PREBUILT_OUTPUTS:
        if dst != target_dir:
            continue
        copied = sync_tree(src, stage_dir)
        if copied:
            restored_any = True
            print(f"[sync] copied {copied} files from {src} -> {stage_dir}")

    for label, config_path, native_root, output_root in EXACT_NATIVE_BUNDLES:
        if output_root != target_dir:
            continue
        if restore_native_bundle(label, config_path, native_root, stage_dir):
            restored_any = True

    if not restored_any:
        shutil.rmtree(stage_dir)
        print(f"[{ANIMATION_SPINE_BUNDLE}] skipped: no source data was found")
        return

    backup_dir = backup_existing_dir(target_dir)
    stage_dir.rename(target_dir)
    if backup_dir:
        print(f"[{ANIMATION_SPINE_BUNDLE}] previous output moved to: {display_path(backup_dir)}")


def main() -> None:
    OUTPUT_BASE.mkdir(parents=True, exist_ok=True)
    BACKUP_BASE.mkdir(parents=True, exist_ok=True)
    cleanup_stale_stage_dirs()
    migrate_legacy_backups()
    if not ONLY_BUNDLES or ANIMATION_SPINE_BUNDLE in ONLY_BUNDLES:
        restore_animation_spine()
    restore_local_bundles()
    restore_gamecache_bundles()
    if not ONLY_BUNDLES or SPINE_CACHE_BUNDLE in ONLY_BUNDLES or f"gamecaches/{SPINE_CACHE_BUNDLE}" in ONLY_BUNDLES:
        restore_gamecache_spine()
    print("Done. Restored assets are under: pack/restored")


if __name__ == "__main__":
    main()
