#!/usr/bin/env python3
"""Download Kenney Space Kit (CC0) → assets/ (GLTF or OBJ)."""

from __future__ import annotations

import json
import shutil
import urllib.request
import zipfile
from pathlib import Path

URL = "https://kenney.nl/media/pages/assets/space-kit/cceeafbd0c-1677698978/kenney_space-kit.zip"
ROOT = Path(__file__).resolve().parent.parent
ASSETS = ROOT / "assets"
ZIP_PATH = Path("/tmp/kenney-space-kit.zip")

SHIP_KEYS = ("craft_racer", "craft_speeder", "craft_miner", "craft_cargo", "ship")
ASTEROID_KEYS = ("meteor_detailed", "meteor_half", "meteor", "rock")
PROJECTILE_KEYS = ("laser", "rocket", "weapon", "bullet")


def download_zip() -> zipfile.ZipFile:
    print("Downloading Kenney Space Kit (CC0)…")
    urllib.request.urlretrieve(URL, ZIP_PATH)
    return zipfile.ZipFile(ZIP_PATH)


def index_members(zf: zipfile.ZipFile) -> dict[str, list[str]]:
    by_stem: dict[str, list[str]] = {}
    for name in zf.namelist():
        if name.endswith("/"):
            continue
        stem = Path(name).stem.lower()
        by_stem.setdefault(stem, []).append(name)
    return by_stem


def pick_path(by_stem: dict[str, list[str]], keys: tuple[str, ...], ext: str) -> str | None:
    for key in keys:
        kl = key.lower()
        for stem, paths in by_stem.items():
            if kl in stem:
                for p in paths:
                    if p.lower().endswith(ext):
                        return p
    return None


def patch_gltf_bin_uri(gltf_path: Path, bin_name: str) -> None:
    data = json.loads(gltf_path.read_text(encoding="utf-8"))
    for buf in data.get("buffers", []):
        if "uri" in buf and str(buf["uri"]).endswith(".bin"):
            buf["uri"] = bin_name
    gltf_path.write_text(json.dumps(data, separators=(",", ":")), encoding="utf-8")


def extract_gltf(zf: zipfile.ZipFile, member: str, dest_stem: str) -> None:
    ASSETS.mkdir(parents=True, exist_ok=True)
    tmp = ASSETS / "_tmp"
    tmp.mkdir(exist_ok=True)
    zf.extract(member, tmp)
    src = tmp / member
    dest_gltf = ASSETS / f"{dest_stem}.gltf"
    shutil.move(src, dest_gltf)

    bin_member = member.replace(".gltf", ".bin")
    if bin_member in zf.namelist():
        zf.extract(bin_member, tmp)
        src_bin = tmp / bin_member
        dest_bin = ASSETS / f"{dest_stem}.bin"
        shutil.move(src_bin, dest_bin)
        patch_gltf_bin_uri(dest_gltf, f"{dest_stem}.bin")

    shutil.rmtree(tmp, ignore_errors=True)
    print(f"  glTF → {dest_stem}.gltf")


def extract_obj_mtl(zf: zipfile.ZipFile, obj_member: str, dest_stem: str) -> None:
    ASSETS.mkdir(parents=True, exist_ok=True)
    tmp = ASSETS / "_tmp"
    tmp.mkdir(exist_ok=True)

    zf.extract(obj_member, tmp)
    shutil.move(tmp / obj_member, ASSETS / f"{dest_stem}.obj")

    mtl_member = obj_member.replace(".obj", ".mtl")
    if mtl_member in zf.namelist():
        zf.extract(mtl_member, tmp)
        shutil.move(tmp / mtl_member, ASSETS / f"{dest_stem}.mtl")

    shutil.rmtree(tmp, ignore_errors=True)
    print(f"  OBJ  → {dest_stem}.obj")


def main() -> None:
    zf = download_zip()
    by_stem = index_members(zf)

    exts = sorted({Path(n).suffix.lower() for names in by_stem.values() for n in names})
    print(f"Archive extensions: {', '.join(exts)}")

    use_gltf = any(p.endswith(".gltf") for names in by_stem.values() for p in names)
    use_obj = any(p.endswith(".obj") for names in by_stem.values() for p in names)

    if not use_gltf and not use_obj:
        sample = [n for names in by_stem.values() for n in names][:30]
        print("Sample paths:", *sample, sep="\n  ")
        raise SystemExit("No GLTF or OBJ models in archive")

    fmt = "gltf" if use_gltf else "obj"
    ext = f".{fmt}" if fmt == "gltf" else ".obj"
    print(f"Using format: {fmt.upper()}")

    ship = pick_path(by_stem, SHIP_KEYS, ext)
    asteroid = pick_path(by_stem, ASTEROID_KEYS, ext)
    projectile = pick_path(by_stem, PROJECTILE_KEYS, ext)

    if not ship:
        print("Craft-like stems:", [s for s in by_stem if "craft" in s][:15])
        raise SystemExit("Ship model not found")

    print("Extracting…")
    if use_gltf:
        extract_gltf(zf, ship, "ship")
        if asteroid:
            extract_gltf(zf, asteroid, "asteroid")
        else:
            shutil.copy(ASSETS / "ship.gltf", ASSETS / "asteroid.gltf")
            if (ASSETS / "ship.bin").exists():
                shutil.copy(ASSETS / "ship.bin", ASSETS / "asteroid.bin")
                patch_gltf_bin_uri(ASSETS / "asteroid.gltf", "asteroid.bin")
        if projectile:
            extract_gltf(zf, projectile, "projectile")
        else:
            shutil.copy(ASSETS / "ship.gltf", ASSETS / "projectile.gltf")
            if (ASSETS / "ship.bin").exists():
                shutil.copy(ASSETS / "ship.bin", ASSETS / "projectile.bin")
                patch_gltf_bin_uri(ASSETS / "projectile.gltf", "projectile.bin")
    else:
        extract_obj_mtl(zf, ship, "ship")
        if asteroid:
            extract_obj_mtl(zf, asteroid, "asteroid")
        else:
            shutil.copy(ASSETS / "ship.obj", ASSETS / "asteroid.obj")
            if (ASSETS / "ship.mtl").exists():
                shutil.copy(ASSETS / "ship.mtl", ASSETS / "asteroid.mtl")
        if projectile:
            extract_obj_mtl(zf, projectile, "projectile")
        else:
            shutil.copy(ASSETS / "ship.obj", ASSETS / "projectile.obj")
            if (ASSETS / "ship.mtl").exists():
                shutil.copy(ASSETS / "ship.mtl", ASSETS / "projectile.mtl")

    zf.close()
    (ASSETS / "format.txt").write_text(fmt, encoding="utf-8")

    for p in sorted(ASSETS.iterdir()):
        if p.name not in {".gitkeep", "_tmp"}:
            print(f"  {p.name} ({p.stat().st_size} bytes)")
    print(f"Done — {ASSETS}")


if __name__ == "__main__":
    main()