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

"""Output data models matching the existing seed JSON format."""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class Encounter:
    pokeapi_id: int
    pokemon_name: str
    method: str
    encounter_rate: int
    min_level: int
    max_level: int
    conditions: dict[str, int] | None = None

    def to_dict(self) -> dict:
        d: dict = {
            "pokeapi_id": self.pokeapi_id,
            "pokemon_name": self.pokemon_name,
            "method": self.method,
            "min_level": self.min_level,
            "max_level": self.max_level,
        }
        if self.conditions:
            d["encounter_rate"] = None
            d["conditions"] = self.conditions
        else:
            d["encounter_rate"] = self.encounter_rate
        return d


@dataclass
class Route:
    name: str
    order: int
    encounters: list[Encounter] = field(default_factory=list)
    children: list[Route] = field(default_factory=list)

    def to_dict(self) -> dict:
        d: dict = {
            "name": self.name,
            "order": self.order,
            "encounters": [e.to_dict() for e in self.encounters],
        }
        if self.children:
            d["children"] = [c.to_dict() for c in self.children]
        return d


@dataclass
class Game:
    name: str
    slug: str
    generation: int
    region: str
    release_year: int
    color: str | None = None

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "slug": self.slug,
            "generation": self.generation,
            "region": self.region,
            "release_year": self.release_year,
            "color": self.color,
        }


@dataclass
class Pokemon:
    pokeapi_id: int
    national_dex: int
    name: str
    types: list[str]
    sprite_url: str

    def to_dict(self) -> dict:
        return {
            "pokeapi_id": self.pokeapi_id,
            "national_dex": self.national_dex,
            "name": self.name,
            "types": self.types,
            "sprite_url": self.sprite_url,
        }