"""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, )