#!/usr/bin/env python3
"""Build a more useful Godot editor reconstruction from exported assets.

This is a static reconstruction. It does not run target AArch64 binaries.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path


PROJECT_GODOT = """; Engine configuration file.
; Reconstructed from /games/pokemon_pro/assets/godot/main.pck project.binary.
; This opens an exported Godot 4.4.1 resource set in the editor. It is not the
; original source checkout: scenes remain exported .scn resources and scripts
; remain compiled .gdc bytecode with .gd.remap files.

config_version=5

[application]

config/name="pokemon_battle_previs"
run/main_scene="res://scenes/main/main.tscn"
config/features=PackedStringArray("4.4", "Forward Plus")

[autoload]

PokemonData="*res://scripts/pokemon_data.gd"
Spike="*res://scripts/godot_spike_bridge.gd"
AsyncLoadUtilV2="*res://scripts/async_load_util_v2.gd"
SpikeSoundIds="*res://scripts/spike_game_sound_ids.gd"

[display]

window/size/viewport_width=1920
window/size/viewport_height=1080
window/stretch/mode="canvas_items"

[gui]

theme/custom_font="uid://mi8jb2bietkw"

[rendering]

renderer/rendering_method="forward_plus"
"""


README = """# Pokemon Spike 3 Godot Editor Reconstruction

This directory is a derived, editor-openable reconstruction from the exported
Godot pack at `/games/pokemon_pro/assets/godot/main.pck`.

Compared with the first wrapper reconstruction, this version also restores
source-like PNG files at the original `res://.../*.png` paths by carving the
embedded WebP/JPEG payloads from Godot `.ctex` cache files and converting them
to PNG. It also includes loose non-Godot assets from `/games/pokemon_pro/assets`
under `rootfs_loose_assets/`.

Important limitations:

- This is not the original Godot source checkout.
- GDScript source files are not present; the pack contains compiled `.gdc`
  bytecode and `.gd.remap` files.
- Scenes are exported binary `.scn` resources reached through `.tscn.remap`
  files, not original editable text `.tscn` source.
- Loose Radium assets and rootfs media are included for inspection, but they
  are not native Godot editor resources.
- No standalone Godot audio resources were found in `main.pck`.

Primary evidence:

- Source pack: `/games/pokemon_pro/assets/godot/main.pck`
- Source pack SHA-256:
  `04af0fd78866ec3ae60b10311dedb3df2f3877595d19ee7bb01665ab10512778`
- Runtime string: `Godot Engine v4.4.1.stable.custom_build`
- Main scene: `res://scenes/main/main.tscn`
"""


IMPORT_PATH_RE = re.compile(r'^path="res://(?P<path>[^"]+)"')


def ignore_metadata(_dir: str, names: list[str]) -> set[str]:
    return {name for name in names if name in {".DS_Store", "__MACOSX"}}


def copy_tree(src: Path, dst: Path) -> None:
    shutil.copytree(src, dst, ignore=ignore_metadata)


def extract_riff_webp(data: bytes) -> bytes | None:
    idx = 0
    while True:
        idx = data.find(b"RIFF", idx)
        if idx < 0:
            return None
        if idx + 12 <= len(data) and data[idx + 8 : idx + 12] == b"WEBP":
            size = int.from_bytes(data[idx + 4 : idx + 8], "little") + 8
            end = idx + size
            if end <= len(data):
                return data[idx:end]
        idx += 4


def extract_jpeg(data: bytes) -> bytes | None:
    start = data.find(b"\xff\xd8\xff")
    if start < 0:
        return None
    end = data.find(b"\xff\xd9", start + 2)
    if end < 0:
        return None
    return data[start : end + 2]


def imported_resource_path(import_file: Path) -> str | None:
    for line in import_file.read_text(encoding="utf-8", errors="replace").splitlines():
        match = IMPORT_PATH_RE.match(line)
        if match:
            return match.group("path")
    return None


def convert_image(magick: str, payload: bytes, suffix: str, out_path: Path) -> bool:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(payload)
        tmp_path = Path(tmp.name)
    try:
        result = subprocess.run(
            [magick, str(tmp_path), str(out_path)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=False,
        )
        return result.returncode == 0 and out_path.exists() and out_path.stat().st_size > 0
    finally:
        try:
            tmp_path.unlink()
        except FileNotFoundError:
            pass


def restore_png_sources(project: Path, magick: str) -> dict[str, object]:
    stats: dict[str, object] = {
        "png_import_files": 0,
        "png_sources_restored": 0,
        "webp_payloads": 0,
        "jpeg_payloads": 0,
        "missing_import_targets": [],
        "conversion_failures": [],
    }

    for import_file in sorted(project.rglob("*.png.import")):
        if "/rootfs_loose_assets/" in str(import_file):
            continue
        stats["png_import_files"] = int(stats["png_import_files"]) + 1
        rel_ctex = imported_resource_path(import_file)
        if not rel_ctex:
            stats["missing_import_targets"].append(str(import_file.relative_to(project)))  # type: ignore[index]
            continue
        ctex_path = project / rel_ctex
        if not ctex_path.exists():
            stats["missing_import_targets"].append(str(ctex_path.relative_to(project)))  # type: ignore[index]
            continue
        data = ctex_path.read_bytes()
        webp = extract_riff_webp(data)
        out_png = import_file.with_suffix("")
        if webp is not None:
            stats["webp_payloads"] = int(stats["webp_payloads"]) + 1
            if convert_image(magick, webp, ".webp", out_png):
                stats["png_sources_restored"] = int(stats["png_sources_restored"]) + 1
            else:
                stats["conversion_failures"].append(str(import_file.relative_to(project)))  # type: ignore[index]
            continue
        jpeg = extract_jpeg(data)
        if jpeg is not None:
            stats["jpeg_payloads"] = int(stats["jpeg_payloads"]) + 1
            if convert_image(magick, jpeg, ".jpg", out_png):
                stats["png_sources_restored"] = int(stats["png_sources_restored"]) + 1
            else:
                stats["conversion_failures"].append(str(import_file.relative_to(project)))  # type: ignore[index]
            continue
        stats["conversion_failures"].append(str(import_file.relative_to(project)))  # type: ignore[index]

    return stats


def ffprobe_streams(ffprobe: str, path: Path) -> list[dict[str, str]]:
    result = subprocess.run(
        [
            ffprobe,
            "-v",
            "error",
            "-show_entries",
            "stream=index,codec_type,codec_name",
            "-of",
            "json",
            str(path),
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if result.returncode != 0:
        return []
    try:
        data = json.loads(result.stdout.decode("utf-8", "replace"))
    except json.JSONDecodeError:
        return []
    return data.get("streams", [])


def media_candidates(paths: list[Path]) -> list[Path]:
    candidates: list[Path] = []
    for root in paths:
        for path in root.rglob("*"):
            if not path.is_file():
                continue
            lower = path.name.lower()
            if lower.endswith((".mp4", ".mov", ".m4a", ".asset")):
                candidates.append(path)
    return sorted(candidates)


def scan_media_audio(ffprobe: str, roots: list[Path], out_tsv: Path) -> dict[str, object]:
    rows: list[tuple[str, str, str]] = []
    audio_files = 0
    video_only_files = 0
    probed_files = 0
    for path in media_candidates(roots):
        streams = ffprobe_streams(ffprobe, path)
        if not streams:
            continue
        probed_files += 1
        has_audio = any(stream.get("codec_type") == "audio" for stream in streams)
        has_video = any(stream.get("codec_type") == "video" for stream in streams)
        if has_audio:
            audio_files += 1
        elif has_video:
            video_only_files += 1
        stream_desc = ",".join(
            f"{stream.get('index')}:{stream.get('codec_type')}:{stream.get('codec_name')}"
            for stream in streams
        )
        rel = path.as_posix()
        rows.append((rel, "yes" if has_audio else "no", stream_desc))
    out_tsv.parent.mkdir(parents=True, exist_ok=True)
    out_tsv.write_text(
        "path\thas_audio\tstreams\n"
        + "\n".join("\t".join(row) for row in rows)
        + ("\n" if rows else ""),
        encoding="utf-8",
    )
    return {
        "candidate_files": len(media_candidates(roots)),
        "probed_files": probed_files,
        "files_with_audio_streams": audio_files,
        "video_only_files": video_only_files,
        "scan_output": str(out_tsv),
    }


def zip_dir(src: Path, zip_path: Path) -> dict[str, int]:
    file_count = 0
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as zf:
        for path in sorted(src.rglob("*")):
            if not path.is_file():
                continue
            if path.name == ".DS_Store" or "__MACOSX" in path.parts:
                continue
            file_count += 1
            zf.write(path, path.relative_to(src))
    return {"files": file_count, "bytes": zip_path.stat().st_size}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--godot-raw", type=Path, default=Path("assets/godot_raw"))
    parser.add_argument("--loose-assets", type=Path, default=Path("games/pokemon_pro/assets"))
    parser.add_argument("--output", type=Path, default=Path("godot_editor_project_reconstructed"))
    parser.add_argument("--zip", type=Path, default=Path("project_reconstructed.zip"))
    parser.add_argument("--media-scan", type=Path, default=Path("analysis/media-audio-scan.tsv"))
    parser.add_argument("--magick", default=shutil.which("magick") or "/opt/homebrew/bin/magick")
    parser.add_argument("--ffprobe", default=shutil.which("ffprobe") or "/opt/homebrew/bin/ffprobe")
    args = parser.parse_args()

    if args.output.exists():
        raise SystemExit(f"output path already exists: {args.output}")
    if args.zip.exists():
        raise SystemExit(f"zip path already exists: {args.zip}")
    if not Path(args.magick).exists():
        raise SystemExit("ImageMagick 'magick' is required for PNG restoration")

    copy_tree(args.godot_raw, args.output)
    (args.output / "project.godot").write_text(PROJECT_GODOT, encoding="utf-8")
    (args.output / "README_RECONSTRUCTION.md").write_text(README, encoding="utf-8")

    restore_stats = restore_png_sources(args.output, args.magick)

    loose_out = args.output / "rootfs_loose_assets"
    for subdir in ("lcd", "nuk"):
        src = args.loose_assets / subdir
        if src.exists():
            copy_tree(src, loose_out / subdir)

    media_stats: dict[str, object] = {}
    if Path(args.ffprobe).exists():
        media_stats = scan_media_audio(args.ffprobe, [args.godot_raw, args.loose_assets], args.media_scan)

    manifest = {
        "source_godot_raw": str(args.godot_raw),
        "source_loose_assets": str(args.loose_assets),
        "output": str(args.output),
        "zip": str(args.zip),
        "png_restore": restore_stats,
        "media_audio_scan": media_stats,
    }
    (args.output / "RECONSTRUCTION_MANIFEST.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    zip_stats = zip_dir(args.output, args.zip)
    manifest["zip_stats"] = zip_stats
    (args.output / "RECONSTRUCTION_MANIFEST.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    print(json.dumps(manifest, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
