from collections import deque

from app.models.evolution import Evolution


def resolve_base_form(pokemon_id: int, evolutions: list[Evolution]) -> int:
    """Walk backward through evolution edges to find the base form of a Pokemon."""
    # Build reverse map: to_pokemon_id -> from_pokemon_id
    reverse: dict[int, list[int]] = {}
    for evo in evolutions:
        reverse.setdefault(evo.to_pokemon_id, []).append(evo.from_pokemon_id)

    visited: set[int] = set()
    current = pokemon_id
    while current not in visited:
        visited.add(current)
        predecessors = reverse.get(current)
        if not predecessors:
            break
        # Follow the first predecessor (handles Shedinja -> Nincada correctly
        # since Nincada evolves *to* both Ninjask and Shedinja)
        current = predecessors[0]
    return current


def build_families(evolutions: list[Evolution]) -> dict[int, list[int]]:
    """Build pokemon_id -> family members mapping using BFS on evolution graph."""
    adj: dict[int, set[int]] = {}
    for evo in evolutions:
        adj.setdefault(evo.from_pokemon_id, set()).add(evo.to_pokemon_id)
        adj.setdefault(evo.to_pokemon_id, set()).add(evo.from_pokemon_id)

    visited: set[int] = set()
    pokemon_to_family: dict[int, list[int]] = {}
    for node in adj:
        if node in visited:
            continue
        component: list[int] = []
        queue = deque([node])
        while queue:
            current = queue.popleft()
            if current in visited:
                continue
            visited.add(current)
            component.append(current)
            for neighbor in adj.get(current, set()):
                if neighbor not in visited:
                    queue.append(neighbor)
        for member in component:
            pokemon_to_family[member] = component
    return pokemon_to_family
