# update_pack_sizes.py
#
# Updates the selected pack by:
# - recalculating each character's sizeBytes from its image file
# - updating registry.json totalSizeBytes
# - updating registry.json characterCount
# - updating registry.json updatedAt
#
# Notes:
# - Missing images are skipped and assigned sizeBytes = 0
# - Missing images are NOT counted in characterCount
# - Character objects are written on a single line for readability

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Optional

REGISTRY_FILENAME = "registry.json"
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".avif")


def format_bytes(size: int) -> str:
    return f"{size / 1024:.1f} KB"


def load_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def save_json(path: Path, data: dict) -> None:
    path.write_text(
        json.dumps(data, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def find_character_image(characters_dir: Path, slug: str) -> Optional[Path]:
    for ext in IMAGE_EXTENSIONS:
        candidate = characters_dir / f"{slug}{ext}"
        if candidate.exists():
            return candidate
    return None


def ordered_character_payload(character: dict) -> dict:
    ordered: dict = {}
    for key in ("name", "slug", "dlc", "sizeBytes"):
        if key in character:
            ordered[key] = character[key]

    for key, value in character.items():
        if key not in ordered:
            ordered[key] = value

    return ordered


def write_manifest(manifest_path: Path, manifest: dict) -> None:
    manifest_copy = dict(manifest)
    characters = manifest_copy.pop("characters", [])

    base_json = json.dumps(manifest_copy, indent=2, ensure_ascii=False)
    if base_json.endswith("\n}"):
        base_json = base_json[:-2]
    elif base_json.endswith("}"):
        base_json = base_json[:-1]
    else:
        raise ValueError("Unexpected JSON format while writing manifest.")

    if characters:
        character_lines = [
            "    " + json.dumps(ordered_character_payload(character), ensure_ascii=False)
            for character in characters
        ]
        characters_json = (
            '  "characters": [\n'
            + ",\n".join(character_lines)
            + "\n  ]\n"
            + "}"
        )
    else:
        characters_json = '  "characters": []\n}'

    manifest_path.write_text(
        base_json + ",\n" + characters_json + "\n",
        encoding="utf-8",
    )


def select_pack(packs: list[dict]) -> Optional[dict]:
    print("\nAvailable packs:\n")
    for index, pack in enumerate(packs, start=1):
        print(f"{index}. {pack['name']} ({pack['id']})")

    raw = input("\nSelect pack number: ").strip()
    try:
        selected_index = int(raw) - 1
    except ValueError:
        print("\n❌ Invalid number")
        return None

    if selected_index < 0 or selected_index >= len(packs):
        print("\n❌ Invalid selection")
        return None

    return packs[selected_index]


def update_pack(root: Path, registry: dict, selected_pack: dict) -> None:
    pack_dir = root / selected_pack["id"]
    manifest_path = pack_dir / "manifest.json"
    characters_dir = pack_dir / "characters"

    if not manifest_path.exists():
        print("\n❌ manifest.json not found")
        return

    if not characters_dir.exists():
        print("\n❌ characters folder not found")
        return

    manifest = load_json(manifest_path)
    characters = manifest.get("characters", [])

    print(f'\nUpdating "{manifest["name"]}"...\n')

    total_size_bytes = 0
    valid_character_count = 0

    for character in characters:
        slug = character["slug"]
        image_path = find_character_image(characters_dir, slug)

        if image_path is None:
            character["sizeBytes"] = 0
            print(f"⚠ Missing image for {slug}")
            continue

        size = image_path.stat().st_size
        character["sizeBytes"] = size
        total_size_bytes += size
        valid_character_count += 1

        print(f"✓ {slug} -> {format_bytes(size)}")

    registry_pack = next(
        (pack for pack in registry.get("packs", []) if pack["id"] == manifest["id"]),
        None,
    )

    if registry_pack is None:
        print("\n❌ Pack not found in registry")
        return

    registry_pack["totalSizeBytes"] = total_size_bytes
    registry_pack["characterCount"] = valid_character_count
    registry["updatedAt"] = datetime.now().strftime("%Y-%m-%d")

    write_manifest(manifest_path, manifest)
    save_json(root / REGISTRY_FILENAME, registry)

    print("\n✅ Done!\n")
    print(f"Characters: {valid_character_count}")
    print(f"Total size: {format_bytes(total_size_bytes)} ({total_size_bytes} bytes)")


def main() -> None:
    root = Path(__file__).parent.resolve()
    registry_path = root / REGISTRY_FILENAME

    if not registry_path.exists():
        print("❌ registry.json not found")
        return

    registry = load_json(registry_path)
    packs = registry.get("packs", [])

    if not packs:
        print("❌ No packs found in registry")
        return

    selected_pack = select_pack(packs)
    if selected_pack is None:
        return

    update_pack(root, registry, selected_pack)


if __name__ == "__main__":
    main()