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

"""Load and validate PokeDB JSON export files."""

from __future__ import annotations

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

_CDN_BASE = "https://cdn.pokedb.org"

REQUIRED_FILES: dict[str, str] = {
    "encounters.json": f"{_CDN_BASE}/data_export_encounters_json",
    "locations.json": f"{_CDN_BASE}/data_export_locations_json",
    "location_areas.json": f"{_CDN_BASE}/data_export_location_areas_json",
    "encounter_methods.json": f"{_CDN_BASE}/data_export_encounter_methods_json",
    "versions.json": f"{_CDN_BASE}/data_export_versions_json",
    "pokemon_forms.json": f"{_CDN_BASE}/data_export_pokemon_forms_json",
}


class PokeDBData:
    """Container for all loaded PokeDB export data."""

    def __init__(
        self,
        encounters: list[dict[str, Any]],
        locations: list[dict[str, Any]],
        location_areas: list[dict[str, Any]],
        encounter_methods: list[dict[str, Any]],
        versions: list[dict[str, Any]],
        pokemon_forms: list[dict[str, Any]],
    ) -> None:
        self.encounters = encounters
        self.locations = locations
        self.location_areas = location_areas
        self.encounter_methods = encounter_methods
        self.versions = versions
        self.pokemon_forms = pokemon_forms

    def summary(self) -> str:
        return (
            f"PokeDB data loaded:\n"
            f"  encounters:        {len(self.encounters):,}\n"
            f"  locations:         {len(self.locations):,}\n"
            f"  location_areas:    {len(self.location_areas):,}\n"
            f"  encounter_methods: {len(self.encounter_methods):,}\n"
            f"  versions:          {len(self.versions):,}\n"
            f"  pokemon_forms:     {len(self.pokemon_forms):,}"
        )


def download_pokedb_data(data_dir: Path) -> None:
    """Download missing PokeDB JSON export files into data_dir."""
    data_dir.mkdir(parents=True, exist_ok=True)
    missing = {f: url for f, url in REQUIRED_FILES.items() if not (data_dir / f).exists()}
    if not missing:
        return

    print(f"Downloading {len(missing)} PokeDB file(s) to {data_dir}...")
    for filename, url in missing.items():
        dest = data_dir / filename
        print(f"  {filename}...", end="", flush=True)
        try:
            urllib.request.urlretrieve(url, dest)
            size_mb = dest.stat().st_size / (1024 * 1024)
            print(f" {size_mb:.1f} MB")
        except Exception as e:
            print(f" FAILED", file=sys.stderr)
            print(f"Error downloading {url}: {e}", file=sys.stderr)
            sys.exit(1)
    print()


def load_pokedb_data(data_dir: Path) -> PokeDBData:
    """Load all PokeDB JSON export files from a directory.

    Downloads any missing files automatically, then validates and loads them.
    """
    download_pokedb_data(data_dir)

    missing = [f for f in REQUIRED_FILES if not (data_dir / f).exists()]
    if missing:
        print(f"Error: Still missing files after download: {missing}", file=sys.stderr)
        sys.exit(1)

    def _load(filename: str) -> list[dict[str, Any]]:
        path = data_dir / filename
        try:
            with open(path) as f:
                data = json.load(f)
        except json.JSONDecodeError as e:
            print(f"Error: Failed to parse {path}: {e}", file=sys.stderr)
            sys.exit(1)

        if not isinstance(data, list):
            print(
                f"Error: Expected a JSON array in {path}, got {type(data).__name__}",
                file=sys.stderr,
            )
            sys.exit(1)

        return data

    return PokeDBData(
        encounters=_load("encounters.json"),
        locations=_load("locations.json"),
        location_areas=_load("location_areas.json"),
        encounter_methods=_load("encounter_methods.json"),
        versions=_load("versions.json"),
        pokemon_forms=_load("pokemon_forms.json"),
    )


class SeedConfig:
    """Container for existing seed configuration files."""

    def __init__(
        self,
        version_groups: dict[str, Any],
        route_order: dict[str, list[str]],
        special_encounters: dict[str, Any] | None,
    ) -> None:
        self.version_groups = version_groups
        self.route_order = route_order
        self.special_encounters = special_encounters


def load_seed_config(seeds_dir: Path) -> SeedConfig:
    """Load existing seed configuration files (version_groups, route_order, etc.).

    Exits with an error message if required config files are missing.
    """
    vg_path = seeds_dir / "version_groups.json"
    if not vg_path.exists():
        print(f"Error: version_groups.json not found at {vg_path}", file=sys.stderr)
        sys.exit(1)

    with open(vg_path) as f:
        version_groups = json.load(f)

    # Load route_order.json and resolve aliases
    ro_path = seeds_dir / "route_order.json"
    if not ro_path.exists():
        print(f"Error: route_order.json not found at {ro_path}", file=sys.stderr)
        sys.exit(1)

    with open(ro_path) as f:
        ro_raw = json.load(f)

    route_order: dict[str, list[str]] = dict(ro_raw.get("routes", {}))

    # Load special_encounters.json (optional)
    se_path = seeds_dir / "special_encounters.json"
    special_encounters = None
    if se_path.exists():
        with open(se_path) as f:
            special_encounters = json.load(f)

    return SeedConfig(
        version_groups=version_groups,
        route_order=route_order,
        special_encounters=special_encounters,
    )