nuzlocke-tracker/scripts/fetch_boss_sprites.py

#!/usr/bin/env python3
"""One-time script to fetch boss battle sprites from Bulbapedia archives.

For trainer bosses (gym leaders, elite four, champions, kahunas):
  Downloads VS portraits or battle sprites from archives.bulbagarden.net.

For totem/noble pokemon bosses:
  Links to existing pokemon sprites already in the project.

Usage:
    python scripts/fetch_boss_sprites.py
"""

import json
import re
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SEED_DIR = ROOT / "backend" / "src" / "app" / "seeds" / "data"
SPRITE_DIR = ROOT / "frontend" / "public" / "boss-sprites"

BULBA_API = "https://archives.bulbagarden.net/w/api.php"
USER_AGENT = "nuzlocke-tracker-sprite-fetch/1.0"

# ── Game slug → Bulbapedia sprite naming conventions ──
# For Gen 1-5: Spr_{CODE}_{Name}.png (pixel battle sprites)
# For Gen 6+:  VS{Name}.png or VS{Name}_{CODE}.png (VS portraits)
GAME_SPRITE_CONFIG: dict[str, dict] = {
    "red":                {"prefix": "Spr", "code": "RG",   "fmt": "spr"},
    "yellow":             {"prefix": "Spr", "code": "Y",    "fmt": "spr"},
    "gold":               {"prefix": "Spr", "code": "GS",   "fmt": "spr"},
    "crystal":            {"prefix": "Spr", "code": "C",    "fmt": "spr"},
    "ruby":               {"prefix": "Spr", "code": "RS",   "fmt": "spr"},
    "emerald":            {"prefix": "Spr", "code": "E",    "fmt": "spr"},
    "firered":            {"prefix": "Spr", "code": "FRLG", "fmt": "spr"},
    "diamond":            {"prefix": "Spr", "code": "DP",   "fmt": "spr"},
    "platinum":           {"prefix": "Spr", "code": "Pt",   "fmt": "spr"},
    "heartgold":          {"prefix": "Spr", "code": "HGSS", "fmt": "spr"},
    "black":              {"prefix": "Spr", "code": "BW",   "fmt": "spr"},
    "black-2":            {"prefix": "Spr", "code": "B2W2", "fmt": "spr"},
    "x":                  {"suffix": "",     "fmt": "vs"},
    "omega-ruby":         {"suffix": "",     "fmt": "vs"},
    "sun":                {"suffix": "",     "fmt": "vs"},
    "ultra-sun":          {"suffix": "USUM", "fmt": "vs"},
    "lets-go-pikachu":    {"suffix": "PE",   "fmt": "vs"},
    "sword":              {"suffix": "",     "fmt": "vs"},
    "brilliant-diamond":  {"suffix": "BDSP", "fmt": "vs"},
    "legends-arceus":     {"suffix": "LA",   "fmt": "vs"},
    "scarlet":            {"suffix": "",     "fmt": "vs"},
}

# ── Boss name → Bulbapedia filename overrides ──
# For names that don't map cleanly to Bulbapedia filenames
NAME_OVERRIDES: dict[str, str] = {
    "Lt. Surge": "Lt_Surge",
    "Crasher Wake": "Crasher_Wake",
    "Professor Kukui": "Professor_Kukui",
    "Tate & Lisa": "Tate_and_Liza",
    "Tate & Liza": "Tate_and_Liza",
    "Top Champion Geeta": "Geeta",
}

# ── Totem/Noble pokemon → pokemon name in seed data for sprite lookup ──
TOTEM_POKEMON: dict[str, str] = {
    "Totem Gumshoos": "Gumshoos",
    "Totem Wishiwashi": "Wishiwashi Solo",
    "Totem Salazzle": "Salazzle",
    "Totem Lurantis": "Lurantis",
    "Totem Vikavolt": "Vikavolt",
    "Totem Mimikyu": "Mimikyu Disguised",
    "Totem Kommo-o": "Kommo O",
    "Totem Araquanid": "Araquanid",
    "Totem Togedemaru": "Togedemaru",
    "Totem Ribombee": "Ribombee",
    # Legends: Arceus nobles
    "Lord Kleavor": "Kleavor",
    "Lady Lilligant": "Lilligant (Hisui)",
    "Lord Arcanine": "Arcanine (Hisui)",
    "Lord Electrode": "Electrode (Hisui)",
    "Lord Avalugg": "Avalugg (Hisui)",
    "Origin Dialga / Palkia": "Dialga (Origin)",
    "Arceus": "Arceus",
    # Scarlet/Violet Titan pokemon
    "Stony Cliff Titan": "Klawf",
    "Open Sky Titan": "Bombirdier",
    "Lurking Steel Titan": "Orthworm",
    "Quaking Earth Titan": "Great Tusk",
    "False Dragon Titan": "Dondozo",
}

# Multi-boss entries: use the first trainer's name
MULTI_BOSS_PRIMARY: dict[str, str] = {
    "Cilan / Chili / Cress": "Cilan",
}


def slugify(name: str) -> str:
    slug = name.lower().replace(" ", "-")
    return re.sub(r"[^a-z0-9-]", "", slug)


def bulba_filename(boss_name: str, config: dict) -> str | None:
    """Construct the Bulbapedia filename for a trainer boss."""
    name = MULTI_BOSS_PRIMARY.get(boss_name, boss_name)
    name = NAME_OVERRIDES.get(name, name)

    # Replace spaces with underscores for Bulbapedia
    bulba_name = name.replace(" ", "_")

    if config["fmt"] == "spr":
        return f"Spr_{config['code']}_{bulba_name}.png"
    else:
        suffix = config.get("suffix", "")
        if suffix:
            return f"VS{bulba_name}_{suffix}.png"
        else:
            return f"VS{bulba_name}.png"


def resolve_image_urls(filenames: list[str]) -> dict[str, str | None]:
    """Batch-resolve Bulbapedia filenames to direct download URLs via the MediaWiki API."""
    result: dict[str, str | None] = {}
    batch_size = 50

    for i in range(0, len(filenames), batch_size):
        batch = filenames[i : i + batch_size]
        titles = "|".join(f"File:{fn}" for fn in batch)
        url = f"{BULBA_API}?action=query&titles={urllib.request.quote(titles)}&prop=imageinfo&iiprop=url&format=json"

        req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = json.loads(resp.read())
        except (urllib.error.URLError, OSError) as exc:
            print(f"  API error: {exc}")
            for fn in batch:
                result[fn] = None
            continue

        pages = data.get("query", {}).get("pages", {})
        # Build title → url mapping
        page_urls: dict[str, str] = {}
        for page in pages.values():
            info = page.get("imageinfo", [])
            if info:
                # Normalize title: "File:Spr BW Cilan.png" → "Spr_BW_Cilan.png"
                title = page["title"].removeprefix("File:").replace(" ", "_")
                page_urls[title] = info[0]["url"]

        for fn in batch:
            result[fn] = page_urls.get(fn)

        if i + batch_size < len(filenames):
            time.sleep(0.5)  # be polite

    return result


def download_image(url: str, dest: Path) -> bool:
    """Download an image file."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            dest.write_bytes(resp.read())
        return True
    except (urllib.error.URLError, OSError) as exc:
        print(f"  FAILED to download {url}: {exc}")
        return False


def load_pokemon_sprites() -> dict[str, str]:
    """Load pokemon name → sprite_url mapping from the pokemon seed data."""
    pokemon_file = SEED_DIR / "pokemon.json"
    if not pokemon_file.exists():
        return {}
    with open(pokemon_file) as f:
        pokemon_list = json.load(f)
    return {p["name"]: p.get("sprite_url", "") for p in pokemon_list if p.get("sprite_url")}


def main():
    pokemon_sprites = load_pokemon_sprites()
    if not pokemon_sprites:
        print("Warning: Could not load pokemon sprites for totem/noble mapping")

    # Collect all filenames we need to resolve
    all_filenames: list[str] = []
    # Track: (game_slug, boss_index, bulba_filename, dest_path, is_pokemon)
    tasks: list[tuple[str, int, str | None, Path, bool]] = []

    boss_files = sorted(SEED_DIR.glob("*-bosses.json"))

    for boss_file in boss_files:
        game_slug = boss_file.name.removesuffix("-bosses.json")
        config = GAME_SPRITE_CONFIG.get(game_slug)
        if config is None:
            print(f"Skipping {game_slug}: no sprite config defined")
            continue

        with open(boss_file) as f:
            bosses = json.load(f)

        sprite_subdir = SPRITE_DIR / game_slug
        for idx, boss in enumerate(bosses):
            boss_name = boss["name"]
            boss_slug = slugify(boss_name)
            dest = sprite_subdir / f"{boss_slug}.png"
            local_path = f"/boss-sprites/{game_slug}/{boss_slug}.png"

            # Check if sprite already exists on disk
            if dest.exists():
                # Ensure seed file has the path
                if boss.get("sprite_url") != local_path:
                    boss["sprite_url"] = local_path
                tasks.append((game_slug, idx, None, dest, False))
                continue

            # Check if this is a totem/noble pokemon boss
            if boss_name in TOTEM_POKEMON:
                pokemon_name = TOTEM_POKEMON[boss_name]
                sprite = pokemon_sprites.get(pokemon_name)
                if sprite:
                    boss["sprite_url"] = sprite
                    tasks.append((game_slug, idx, None, dest, True))
                    continue
                else:
                    print(f"  Warning: No pokemon sprite found for {boss_name} ({pokemon_name})")

            # Construct Bulbapedia filename
            fn = bulba_filename(boss_name, config)
            if fn:
                all_filenames.append(fn)
                tasks.append((game_slug, idx, fn, dest, False))
            else:
                print(f"  Could not construct filename for {boss_name}")
                tasks.append((game_slug, idx, None, dest, False))

        # Write back any sprite_url fixes
        with open(boss_file, "w") as f:
            json.dump(bosses, f, indent=2, ensure_ascii=False)
            f.write("\n")

    if not all_filenames:
        print("No sprites to download!")
        return

    # Deduplicate filenames (some bosses appear across games with same Bulba file)
    unique_filenames = list(dict.fromkeys(all_filenames))
    print(f"\nResolving {len(unique_filenames)} unique Bulbapedia filenames...")
    url_map = resolve_image_urls(unique_filenames)

    resolved = sum(1 for v in url_map.values() if v)
    print(f"Resolved {resolved}/{len(unique_filenames)} URLs")

    # Report missing
    missing = [fn for fn, url in url_map.items() if not url]
    if missing:
        print(f"\nCould not resolve {len(missing)} filenames:")
        for fn in missing:
            print(f"  - {fn}")

    # Download sprites and update seed files
    downloaded = 0
    skipped = 0
    failed = 0

    # Re-read seed files for updating
    seed_data: dict[str, list] = {}
    for boss_file in boss_files:
        game_slug = boss_file.name.removesuffix("-bosses.json")
        with open(boss_file) as f:
            seed_data[game_slug] = json.load(f)

    for game_slug, idx, bulba_fn, dest, is_pokemon in tasks:
        if is_pokemon or dest.exists():
            skipped += 1
            continue

        if bulba_fn is None:
            continue

        url = url_map.get(bulba_fn)
        if not url:
            failed += 1
            continue

        boss_slug = dest.stem
        local_path = f"/boss-sprites/{game_slug}/{boss_slug}.png"

        if download_image(url, dest):
            seed_data[game_slug][idx]["sprite_url"] = local_path
            downloaded += 1
            print(f"  {game_slug}/{boss_slug}.png")
        else:
            failed += 1

        time.sleep(0.3)  # rate limit

    # Write updated seed files
    for game_slug, bosses in seed_data.items():
        boss_file = SEED_DIR / f"{game_slug}-bosses.json"
        with open(boss_file, "w") as f:
            json.dump(bosses, f, indent=2, ensure_ascii=False)
            f.write("\n")

    print(f"\nDone! Downloaded: {downloaded}, Skipped (existing): {skipped}, Failed: {failed}")

    # Fallback attempts for failed sprites
    if missing:
        print("\n--- Attempting fallback filenames for missing sprites ---")
        _attempt_fallbacks(missing, url_map, seed_data, tasks)


def _attempt_fallbacks(
    missing_filenames: list[str],
    url_map: dict[str, str | None],
    seed_data: dict[str, list],
    tasks: list,
):
    """Try alternate Bulbapedia filename patterns for sprites that weren't found."""
    fallback_filenames: list[str] = []
    fallback_map: dict[str, str] = {}  # fallback_fn → original_fn

    for orig_fn in missing_filenames:
        # Try VS ↔ Spr swaps
        if orig_fn.startswith("Spr_"):
            # Spr_CODE_Name.png → VSName.png
            parts = orig_fn.removeprefix("Spr_").removesuffix(".png").split("_", 1)
            if len(parts) == 2:
                name = parts[1]
                alt = f"VS{name}.png"
                fallback_filenames.append(alt)
                fallback_map[alt] = orig_fn
        elif orig_fn.startswith("VS"):
            # VSName.png → VSName_2.png (alternate appearance)
            base = orig_fn.removesuffix(".png")
            alt = f"{base}_2.png"
            fallback_filenames.append(alt)
            fallback_map[alt] = orig_fn

    if not fallback_filenames:
        return

    print(f"Trying {len(fallback_filenames)} fallback filenames...")
    fallback_urls = resolve_image_urls(fallback_filenames)

    found = 0
    for fb_fn, fb_url in fallback_urls.items():
        if fb_url:
            orig_fn = fallback_map[fb_fn]
            url_map[orig_fn] = fb_url
            found += 1

            # Find and download for matching tasks
            for game_slug, idx, bulba_fn, dest, is_pokemon in tasks:
                if bulba_fn == orig_fn and not dest.exists():
                    boss_slug = dest.stem
                    local_path = f"/boss-sprites/{game_slug}/{boss_slug}.png"
                    if download_image(fb_url, dest):
                        seed_data[game_slug][idx]["sprite_url"] = local_path
                        print(f"  {game_slug}/{boss_slug}.png (fallback: {fb_fn})")
                    time.sleep(0.3)

    if found:
        # Re-save seed files
        for game_slug, bosses in seed_data.items():
            boss_file = SEED_DIR / f"{game_slug}-bosses.json"
            with open(boss_file, "w") as f:
                json.dump(bosses, f, indent=2, ensure_ascii=False)
                f.write("\n")

    print(f"Fallback resolved: {found}/{len(fallback_filenames)}")


if __name__ == "__main__":
    main()