nuzlocke-tracker/tools/import-pokedb/import_pokedb/sprites.py

"""Download and manage PokeDB pokemon sprites."""

from __future__ import annotations

import sys
import urllib.request
from pathlib import Path
from typing import Any

from .mappings import PokemonMapper


def download_sprites(
    pokemon_mapper: PokemonMapper,
    encountered_form_ids: set[str],
    sprites_dir: Path,
) -> dict[str, str]:
    """Download sprites for all encountered pokemon forms.

    Returns a mapping of pokemon_form_identifier → local sprite filename.
    Skips already-downloaded sprites.
    """
    sprites_dir.mkdir(parents=True, exist_ok=True)

    to_download: list[tuple[str, str, Path]] = []  # (form_id, url, dest)
    result: dict[str, str] = {}

    for form_id in sorted(encountered_form_ids):
        info = pokemon_mapper.lookup(form_id)
        if info is None:
            continue

        pokeapi_id, _ = info
        sprite_url = pokemon_mapper.get_sprite_url(form_id)
        if not sprite_url:
            continue

        filename = f"{pokeapi_id}.webp"
        dest = sprites_dir / filename
        result[form_id] = filename

        if not dest.exists():
            to_download.append((form_id, sprite_url, dest))

    if not to_download:
        print(f"  Sprites: {len(result)} already cached")
        return result

    print(f"  Downloading {len(to_download)} sprites ({len(result) - len(to_download)} cached)...")

    failed = 0
    for i, (form_id, url, dest) in enumerate(to_download, 1):
        try:
            urllib.request.urlretrieve(url, dest)
        except Exception as e:
            print(f"  Warning: Failed to download sprite for {form_id}: {e}", file=sys.stderr)
            failed += 1
            # Remove the failed entry from results
            result.pop(form_id, None)

        # Progress every 100
        if i % 100 == 0:
            print(f"    {i}/{len(to_download)}...")

    if failed:
        print(f"  Sprites: {len(result)} downloaded, {failed} failed")
    else:
        print(f"  Sprites: {len(result)} total ({len(to_download)} new)")

    return result


def download_all_sprites(
    pokemon_mapper: PokemonMapper,
    sprites_dir: Path,
) -> set[int]:
    """Download sprites for all known pokemon (not just encountered ones).

    Returns a set of pokeapi_ids that have sprites downloaded.
    """
    sprites_dir.mkdir(parents=True, exist_ok=True)

    to_download: list[tuple[int, str, Path]] = []
    have_sprites: set[int] = set()

    for pokeapi_id, (_ndex, _name) in pokemon_mapper.all_pokemon():
        if pokemon_mapper.has_sprite_for_id(pokeapi_id):
            filename = f"{pokeapi_id}.webp"
            dest = sprites_dir / filename
            have_sprites.add(pokeapi_id)

            if not dest.exists():
                # Get the sprite URL via the mapper
                url = pokemon_mapper.get_sprite_url_for_id(pokeapi_id)
                if url:
                    to_download.append((pokeapi_id, url, dest))

    if not to_download:
        print(f"  All sprites: {len(have_sprites)} already cached")
        return have_sprites

    print(f"  Downloading {len(to_download)} additional sprites ({len(have_sprites) - len(to_download)} cached)...")

    failed = 0
    for i, (pid, url, dest) in enumerate(to_download, 1):
        try:
            urllib.request.urlretrieve(url, dest)
        except Exception as e:
            print(f"  Warning: Failed to download sprite for pokemon {pid}: {e}", file=sys.stderr)
            failed += 1
            have_sprites.discard(pid)

        if i % 100 == 0:
            print(f"    {i}/{len(to_download)}...")

    if failed:
        print(f"  All sprites: {len(have_sprites)} downloaded, {failed} failed")
    else:
        print(f"  All sprites: {len(have_sprites)} total ({len(to_download)} new)")

    return have_sprites


def sprite_path_for_pokemon(pokeapi_id: int) -> str:
    """Generate the sprite URL path for use in pokemon.json.

    Returns an absolute path like "/sprites/25.webp" for the frontend
    (files in frontend/public/ are served at the root).
    """
    return f"/sprites/{pokeapi_id}.webp"