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

"""Output seed JSON files in the existing format."""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path
from typing import Any

from .loader import SeedConfig
from .mappings import TYPE_ID_MAP, PokemonMapper
from .models import Encounter, Route
from .sprites import sprite_path_for_pokemon

# Max route name length (matches DB column VARCHAR(100))
_MAX_ROUTE_NAME_LEN = 100

# Pattern for Sw/Sh den area names: "Location (Den X - long description - Common/Rare)"
_DEN_NAME_RE = re.compile(r"^(.+?) \(Den ([A-Z]\d*) - .+ - (Common|Rare)\)$")


def _truncate_route_name(name: str) -> str:
    """Truncate a route name to fit the database column limit.

    Applies smart truncation for known patterns (e.g. Sw/Sh den descriptions).
    """
    if len(name) <= _MAX_ROUTE_NAME_LEN:
        return name

    # Sw/Sh dens: shorten "Location (Den X - long desc - Common/Rare)" → "Location (Den X - Common/Rare)"
    m = _DEN_NAME_RE.match(name)
    if m:
        shortened = f"{m.group(1)} (Den {m.group(2)} - {m.group(3)})"
        if len(shortened) <= _MAX_ROUTE_NAME_LEN:
            return shortened

    # Generic truncation: cut at last space before limit, add ellipsis
    truncated = name[:_MAX_ROUTE_NAME_LEN - 1].rsplit(" ", 1)[0] + "…"
    return truncated


# ---------------------------------------------------------------------------
# Route ordering
# ---------------------------------------------------------------------------

def _get_route_order(config: SeedConfig, vg_key: str) -> list[str] | None:
    """Get the route order list for a version group, resolving aliases."""
    return config.route_order.get(vg_key)


def _route_sort_key(name: str, order_list: list[str]) -> tuple[int, str]:
    """Compute sort key for a route name against an ordered list."""
    for i, ordered_name in enumerate(order_list):
        if name == ordered_name or name.startswith(ordered_name + " ("):
            return (i, name)
    return (len(order_list), name)


def sort_routes(routes: list[Route], config: SeedConfig, vg_key: str) -> list[Route]:
    """Sort routes by game progression order and assign sequential order values."""
    order_list = _get_route_order(config, vg_key)
    if order_list:
        routes.sort(key=lambda r: _route_sort_key(r.name, order_list))
    else:
        print(f"  Warning: No route order for {vg_key}, using default order", file=sys.stderr)

    # Assign sequential order values
    order = 1
    for route in routes:
        route.order = order
        order += 1
        for child in route.children:
            child.order = order
            order += 1

    return routes


# ---------------------------------------------------------------------------
# Special encounters
# ---------------------------------------------------------------------------

def _get_special_encounters(
    config: SeedConfig,
    vg_key: str,
) -> dict[str, list[dict[str, Any]]] | None:
    """Get special encounters for a version group, resolving aliases."""
    if not config.special_encounters:
        return None

    encounters = config.special_encounters.get("encounters", {})
    aliases = config.special_encounters.get("aliases", {})

    if vg_key in encounters:
        return encounters[vg_key]
    if vg_key in aliases:
        alias_target = aliases[vg_key]
        if alias_target in encounters:
            return encounters[alias_target]

    return None


def merge_special_encounters(
    routes: list[Route],
    config: SeedConfig,
    vg_key: str,
    pokemon_mapper: PokemonMapper | None = None,
) -> list[Route]:
    """Merge special encounters (starters, gifts, fossils) into routes."""
    special_data = _get_special_encounters(config, vg_key)
    if not special_data:
        return routes

    # Build lookup: route name → route (including children)
    route_map: dict[str, Route] = {}
    for route in routes:
        route_map[route.name] = route
        for child in route.children:
            route_map[child.name] = child

    for location_name, encounter_dicts in special_data.items():
        encounters = []
        for e in encounter_dicts:
            # Look up proper display name from pokemon mapper if available
            pokemon_name = e["pokemon_name"]
            if pokemon_mapper:
                info = pokemon_mapper.lookup_by_id(e["pokeapi_id"])
                if info:
                    pokemon_name = info[1]

            encounters.append(Encounter(
                pokeapi_id=e["pokeapi_id"],
                pokemon_name=pokemon_name,
                method=e["method"],
                encounter_rate=e["encounter_rate"],
                min_level=e["min_level"],
                max_level=e["max_level"],
            ))

        if location_name in route_map:
            existing = route_map[location_name].encounters
            for enc in encounters:
                # If a starter encounter matches an existing gift encounter,
                # update the method to "starter" instead of adding a duplicate.
                if enc.method == "starter":
                    match = next(
                        (e for e in existing
                         if e.pokeapi_id == enc.pokeapi_id and e.method == "gift"),
                        None,
                    )
                    if match:
                        match.method = "starter"
                        continue
                existing.append(enc)
        else:
            new_route = Route(name=location_name, order=0, encounters=encounters)
            routes.append(new_route)
            route_map[location_name] = new_route

    return routes


# ---------------------------------------------------------------------------
# JSON output
# ---------------------------------------------------------------------------

def write_json(path: Path, data: Any) -> None:
    """Write data as formatted JSON with trailing newline."""
    content = json.dumps(data, indent=2, ensure_ascii=False)
    path.write_text(content + "\n", encoding="utf-8")
    print(f"  -> {path}")


def _deduplicate_names(routes: list[Route]) -> None:
    """Ensure all route names are unique by appending a numeric suffix to duplicates."""
    seen: dict[str, int] = {}

    def _unique(name: str) -> str:
        if name not in seen:
            seen[name] = 1
            return name
        seen[name] += 1
        return f"{name} #{seen[name]}"

    for route in routes:
        route.name = _unique(_truncate_route_name(route.name))
        for child in route.children:
            child.name = _unique(_truncate_route_name(child.name))


def write_game_json(routes: list[Route], output_dir: Path, game_slug: str) -> None:
    """Write a per-game route/encounter JSON file."""
    _deduplicate_names(routes)
    data = [r.to_dict() for r in routes]
    write_json(output_dir / f"{game_slug}.json", data)


_GAME_CATEGORY: dict[str, str] = {
    "red": "original", "blue": "original", "yellow": "enhanced",
    "gold": "original", "silver": "original", "crystal": "enhanced",
    "ruby": "original", "sapphire": "original", "emerald": "enhanced",
    "firered": "remake", "leafgreen": "remake",
    "diamond": "original", "pearl": "original", "platinum": "enhanced",
    "heartgold": "remake", "soulsilver": "remake",
    "black": "original", "white": "original",
    "black-2": "sequel", "white-2": "sequel",
    "x": "original", "y": "original",
    "omega-ruby": "remake", "alpha-sapphire": "remake",
    "sun": "original", "moon": "original",
    "ultra-sun": "enhanced", "ultra-moon": "enhanced",
    "lets-go-pikachu": "remake", "lets-go-eevee": "remake",
    "sword": "original", "shield": "original",
    "brilliant-diamond": "remake", "shining-pearl": "remake",
    "legends-arceus": "spinoff",
    "scarlet": "original", "violet": "original",
    "legends-z-a": "spinoff",
}


def write_games_json(config: SeedConfig, output_dir: Path) -> None:
    """Write games.json from version_groups config."""
    games = []
    for vg_info in config.version_groups.values():
        for game_info in vg_info.get("games", {}).values():
            slug = game_info["slug"]
            games.append({
                "name": game_info["name"],
                "slug": slug,
                "generation": vg_info["generation"],
                "region": vg_info["region"],
                "release_year": game_info["release_year"],
                "color": game_info.get("color"),
                "category": _GAME_CATEGORY.get(slug, "original"),
            })
    write_json(output_dir / "games.json", games)
    print(f"  Wrote {len(games)} games")


def write_pokemon_json(
    pokemon_mapper: PokemonMapper,
    encountered_form_ids: set[str],
    sprite_map: dict[str, str],
    output_dir: Path,
) -> None:
    """Write pokemon.json with all known pokemon.

    Includes all pokemon from the existing pokemon.json (base data),
    enriched with PokeDB types and sprite paths for encountered forms.
    """
    # Build a mapping of pokeapi_id → (form_id, form_data) for encountered forms
    encountered_by_id: dict[int, tuple[str, dict[str, Any] | None]] = {}
    for form_id in encountered_form_ids:
        info = pokemon_mapper.lookup(form_id)
        if info is None:
            continue
        pokeapi_id, _ = info
        if pokeapi_id not in encountered_by_id:
            form_data = pokemon_mapper.get_form_data(form_id)
            encountered_by_id[pokeapi_id] = (form_id, form_data)

    pokemon_list: list[dict[str, Any]] = []

    for pokeapi_id, (ndex, name) in pokemon_mapper.all_pokemon():
        # Enrich with PokeDB data if this pokemon was encountered
        if pokeapi_id in encountered_by_id:
            form_id, form_data = encountered_by_id[pokeapi_id]
            types = _extract_types(form_data) if form_data else []
            sprite_url = sprite_path_for_pokemon(pokeapi_id) if form_id in sprite_map else None
            national_dex = form_data.get("ndex_id", ndex) if form_data else ndex
        else:
            # Not encountered — use existing data, try to find PokeDB form for types
            types = pokemon_mapper.get_types_for_id(pokeapi_id)
            sprite_url = sprite_path_for_pokemon(pokeapi_id) if pokemon_mapper.has_sprite_for_id(pokeapi_id) else None
            national_dex = ndex

        pokemon_list.append({
            "pokeapi_id": pokeapi_id,
            "national_dex": national_dex,
            "name": name,
            "types": types,
            "sprite_url": sprite_url,
        })

    # Sort by pokeapi_id
    pokemon_list.sort(key=lambda p: p["pokeapi_id"])

    write_json(output_dir / "pokemon.json", pokemon_list)
    print(f"  Wrote {len(pokemon_list)} pokemon")


def _extract_types(form_data: dict[str, Any]) -> list[str]:
    """Extract type names from a PokeDB form record."""
    types = []
    type1_id = form_data.get("type_1_id")
    type2_id = form_data.get("type_2_id")
    if type1_id and type1_id in TYPE_ID_MAP:
        types.append(TYPE_ID_MAP[type1_id])
    if type2_id and type2_id in TYPE_ID_MAP:
        types.append(TYPE_ID_MAP[type2_id])
    return types