from collections import deque

from app.models.evolution import Evolution


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