import contextlib
from collections import deque
from collections.abc import Iterable

from dr_eureka.core import GameState, Move
from dr_eureka.solver import DrEurekaSolver, SearchResult, SolutionStep


class BFSSolver(DrEurekaSolver):
    """
    Shortest-path solver for the case where every move has equal cost.

    This is deliberately simple and makes a good baseline implementation.
    """

    def iter_neighbours(self, state: GameState) -> Iterable[tuple[Move, GameState]]:
        tube_count = len(state.tubes)

        for src in range(tube_count):
            tube = state.tubes[src]

            if not tube.base_up:
                with contextlib.suppress(ValueError):
                    yield Move(kind="flip", src=src), state.flip_tube_base_up(src)

            for dst in range(tube_count):
                if src == dst:
                    continue
                with contextlib.suppress(ValueError):
                    yield Move(kind="pour", src=src, dst=dst), state.pour_one_ball(src, dst)

    def solve(self, start: GameState, goal: GameState) -> SearchResult | None:
        start_key = self.state_key(start)
        goal_key = self.state_key(goal)

        if start_key == goal_key:
            return SearchResult(steps=(), visited_states=1)

        queue: deque[GameState] = deque([start])
        visited: set[tuple[object, ...]] = {start_key}
        parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}
        state_by_key: dict[tuple[object, ...], GameState] = {start_key: start}

        while queue:
            current = queue.popleft()

            for move, next_state in self.iter_neighbours(current):
                next_key = self.state_key(next_state)
                if next_key in visited:
                    continue

                current_key = self.state_key(current)
                visited.add(next_key)
                parent[next_key] = (current_key, move)
                state_by_key[next_key] = next_state

                if next_key == goal_key:
                    steps = self._reconstruct_steps(parent, state_by_key, start_key, next_key)
                    return SearchResult(
                        steps=steps,
                        visited_states=len(visited),
                    )

                queue.append(next_state)

        return None

    @staticmethod
    def _reconstruct_steps(
        parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        state_by_key: dict[tuple[object, ...], GameState],
        start_key: tuple[object, ...],
        goal_key: tuple[object, ...],
    ) -> tuple[SolutionStep, ...]:
        path: list[SolutionStep] = []
        current_key = goal_key

        while current_key != start_key:
            previous_key, move = parent[current_key]
            path.append(SolutionStep(move=move, state=state_by_key[current_key]))
            current_key = previous_key

        path.reverse()
        return tuple(path)
