"""
Bi-directional BFS for Dr Eureka.

Searches from both start and goal simultaneously until the two frontiers
meet. Forward expansion uses pour and flip; backward expansion reverses both
(pour by undoing the pour, flip by "unflipping" a base-up tube).
The goal is considered reached when the state matches any of the 8 variants
(three tubes, each upright or base-up: 2^3 = 8 equivalent goal states).
"""

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

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


def _unflip_tube(tube: Tube) -> Tube:
    """Predecessor of a flip: tube goes from base-up back to upright (balls reversed)."""
    if not tube.base_up:
        raise ValueError("Tube is not base-up.")
    return Tube(
        balls=tube.balls[::-1],
        base_up=False,
        capacity=tube.capacity,
    )


def goal_variants(goal: GameState) -> Iterable[GameState]:
    """
    Yield all 8 goal states: same tube contents, each tube either upright or base-up.

    A challenge can be completed with any combination of tube orientations.
    """
    n = len(goal.tubes)
    for bits in range(1 << n):
        tubes: list[Tube] = []
        for i in range(n):
            t = goal.tubes[i]
            want_base_up = bool((bits >> i) & 1)
            if t.base_up == want_base_up:
                tubes.append(t)
            elif want_base_up:
                tubes.append(t.flip_base_up())
            else:
                tubes.append(_unflip_tube(t))
        yield GameState(tuple(tubes))


class BiDirectionalCursorSolver(DrEurekaSolver):
    """
    Bi-directional BFS solver.

    Expands from start (forward) and from all 8 goal variants (backward).
    Backward neighbours reverse pour (undo the pour) and reverse flip
    (unflip a base-up tube to get its predecessor).
    """

    def iter_neighbours(self, state: GameState) -> Iterable[tuple[Move, GameState]]:
        """Forward: all valid pour and flip moves from state."""
        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 iter_backward_neighbours(self, state: GameState) -> Iterable[tuple[Move, GameState]]:
        """
        Backward: states that reach `state` in one forward move.

        Reverses pour (state.pour_one_ball(dst, src)) and flip (unflip each
        base-up tube; that predecessor has the tube upright with balls reversed).
        """
        tube_count = len(state.tubes)
        # Reverse pour: predecessor is state with one pour undone
        for dst in range(tube_count):
            for src in range(tube_count):
                if src == dst:
                    continue
                with contextlib.suppress(ValueError):
                    pred = state.pour_one_ball(dst, src)
                    yield Move(kind="pour", src=src, dst=dst), pred
        # Reverse flip: for each base-up tube, predecessor has it unflipped
        for i in range(tube_count):
            if state.tubes[i].base_up:
                with contextlib.suppress(ValueError):
                    pred_tube = _unflip_tube(state.tubes[i])
                    new_tubes = list(state.tubes)
                    new_tubes[i] = pred_tube
                    pred = GameState(tuple(new_tubes))
                    yield Move(kind="flip", src=i), pred

    def solve(self, start: GameState, goal: GameState) -> SearchResult | None:
        start_key = self.state_key(start)
        goal_keys = {self.state_key(g) for g in goal_variants(goal)}

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

        # Forward: from start
        queue_fwd: deque[GameState] = deque([start])
        visited_fwd: set[tuple[object, ...]] = {start_key}
        parent_fwd: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}
        state_by_key_fwd: dict[tuple[object, ...], GameState] = {start_key: start}

        # Backward: from all 8 goal variants
        queue_back: deque[GameState] = deque()
        visited_back: set[tuple[object, ...]] = set()
        parent_back: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}
        state_by_key_back: dict[tuple[object, ...], GameState] = {}
        for g in goal_variants(goal):
            k = self.state_key(g)
            if k not in visited_back:
                visited_back.add(k)
                state_by_key_back[k] = g
                queue_back.append(g)

        while queue_fwd and queue_back:
            # Expand forward (one level)
            for _ in range(len(queue_fwd)):
                if not queue_fwd:
                    break
                current = queue_fwd.popleft()
                current_key = self.state_key(current)

                for move, next_state in self.iter_neighbours(current):
                    next_key = self.state_key(next_state)
                    if next_key in visited_fwd:
                        continue
                    visited_fwd.add(next_key)
                    parent_fwd[next_key] = (current_key, move)
                    state_by_key_fwd[next_key] = next_state

                    if next_key in visited_back:
                        steps = self._reconstruct_steps_meet(
                            parent_fwd,
                            parent_back,
                            state_by_key_fwd,
                            state_by_key_back,
                            start_key,
                            goal_keys,
                            next_key,
                        )
                        return SearchResult(
                            steps=steps,
                            visited_states=len(visited_fwd) + len(visited_back),
                        )
                    queue_fwd.append(next_state)

            # Expand backward (one level)
            for _ in range(len(queue_back)):
                if not queue_back:
                    break
                current = queue_back.popleft()
                current_key = self.state_key(current)

                for move, pred_state in self.iter_backward_neighbours(current):
                    pred_key = self.state_key(pred_state)
                    if pred_key in visited_back:
                        continue
                    visited_back.add(pred_key)
                    # From pred we apply move to get current (closer to goal)
                    parent_back[pred_key] = (current_key, move)
                    state_by_key_back[pred_key] = pred_state

                    if pred_key in visited_fwd:
                        steps = self._reconstruct_steps_meet(
                            parent_fwd,
                            parent_back,
                            state_by_key_fwd,
                            state_by_key_back,
                            start_key,
                            goal_keys,
                            pred_key,
                        )
                        return SearchResult(
                            steps=steps,
                            visited_states=len(visited_fwd) + len(visited_back),
                        )
                    queue_back.append(pred_state)

        return None

    def _reconstruct_steps_meet(
        self,
        parent_fwd: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        parent_back: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        state_by_key_fwd: dict[tuple[object, ...], GameState],
        state_by_key_back: dict[tuple[object, ...], GameState],
        start_key: tuple[object, ...],
        goal_keys: set[tuple[object, ...]],
        meet_key: tuple[object, ...],
    ) -> tuple[SolutionStep, ...]:
        # Path from start to meet: follow parent_fwd from meet to start, then reverse
        path_fwd: list[SolutionStep] = []
        key = meet_key
        while key != start_key:
            prev_key, move = parent_fwd[key]
            path_fwd.append(SolutionStep(move=move, state=state_by_key_fwd[key]))
            key = prev_key
        path_fwd.reverse()

        # Path from meet to (any) goal: follow parent_back until we hit a goal state
        path_back: list[SolutionStep] = []
        key = meet_key
        while key not in goal_keys:
            next_key, move = parent_back[key]
            state_after = state_by_key_back[next_key]
            path_back.append(SolutionStep(move=move, state=state_after))
            key = next_key

        return tuple(path_fwd) + tuple(path_back)


class BiDirectionalGippitySolver(DrEurekaSolver):
    """
    Bidirectional BFS solver for Dr Eureka.

    Notes:
    - The search is still exact and optimal because every move has unit cost.
    - States are normalised internally by sorting their tubes. That keeps the
      stored representative aligned with `state_key()` and makes stitching the
      two search halves together much less cursed.
    - By default, an upright goal tube is treated as orientation-agnostic at
      the very end: either that upright tube or its flipped base-up version is
      accepted as a terminal goal tube. This yields up to 2^n goal states.
      If you want exact goal orientation instead, override `goal_states()`.
    """

    def normalise_state(self, state: GameState) -> GameState:
        return GameState(state.canonical())

    def goal_states(self, goal: GameState) -> tuple[GameState, ...]:
        """
        Expand a displayed goal into all terminal physical goal states.

        For each upright goal tube, both of these are accepted:
        - upright as given
        - flipped base-up, which reverses the visible ball order

        If a tube is already marked base-up explicitly, it is kept fixed.
        """
        variants: list[tuple[Tube, ...]] = [()]

        for tube in goal.tubes:
            options = (tube,) if tube.base_up else (tube, tube.flip_base_up())

            variants = [(*prefix, option) for prefix in variants for option in options]

        unique: dict[tuple[object, ...], GameState] = {}
        for tubes in variants:
            state = self.normalise_state(GameState(tubes))
            unique[self.state_key(state)] = state

        return tuple(unique.values())

    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):
                    next_state = state.flip_tube_base_up(src)
                    yield Move(kind="flip", src=src), self.normalise_state(next_state)

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

    def iter_predecessors(self, state: GameState) -> Iterable[tuple[Move, GameState]]:
        """
        Generate predecessor states in the reverse graph.

        The returned Move is always the *forward* move that takes the predecessor
        to `state`. That makes reconstruction much cleaner.
        """
        tube_count = len(state.tubes)

        # Reverse of a terminal flip:
        #   predecessor tube is upright with reversed balls
        #   forward move is flip(index)
        for index, tube in enumerate(state.tubes):
            if not tube.base_up:
                continue

            predecessor_tubes = list(state.tubes)
            predecessor_tubes[index] = Tube(
                balls=tube.balls[::-1],
                base_up=False,
                capacity=tube.capacity,
            )
            predecessor = self.normalise_state(GameState(tuple(predecessor_tubes)))
            move = self._infer_forward_move(predecessor, state)
            yield move, predecessor

        # Reverse of a one-ball pour:
        # If current dst tube is upright and non-empty, its top ball may have
        # been the ball just poured in. Move that ball back onto the predecessor
        # source tube.
        for dst in range(tube_count):
            dst_tube = state.tubes[dst]
            if dst_tube.base_up or dst_tube.is_empty:
                continue

            moved_ball = dst_tube.balls[0]
            _, predecessor_dst = dst_tube.remove_top_ball()

            for src in range(tube_count):
                if src == dst:
                    continue

                src_tube = state.tubes[src]
                if src_tube.base_up or src_tube.is_full:
                    continue

                predecessor_tubes = list(state.tubes)
                predecessor_tubes[src] = src_tube.add_ball_on_top(moved_ball)
                predecessor_tubes[dst] = predecessor_dst

                predecessor = self.normalise_state(GameState(tuple(predecessor_tubes)))
                move = self._infer_forward_move(predecessor, state)
                yield move, predecessor

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

        backward_roots = self.goal_states(goal)
        backward_root_keys = {self.state_key(state) for state in backward_roots}

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

        forward_queue: deque[GameState] = deque([start_state])
        backward_queue: deque[GameState] = deque(backward_roots)

        forward_visited: set[tuple[object, ...]] = {start_key}
        backward_visited: set[tuple[object, ...]] = set(backward_root_keys)

        forward_parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}
        backward_parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}

        forward_state_by_key: dict[tuple[object, ...], GameState] = {
            start_key: start_state,
        }
        backward_state_by_key: dict[tuple[object, ...], GameState] = {
            self.state_key(state): state for state in backward_roots
        }

        while forward_queue and backward_queue:
            meeting_key = self._expand_forward_layer(
                queue=forward_queue,
                visited=forward_visited,
                other_visited=backward_visited,
                parent=forward_parent,
                state_by_key=forward_state_by_key,
            )
            if meeting_key is not None:
                return self._build_result(
                    meeting_key=meeting_key,
                    start_key=start_key,
                    backward_root_keys=backward_root_keys,
                    forward_parent=forward_parent,
                    backward_parent=backward_parent,
                    forward_state_by_key=forward_state_by_key,
                    backward_state_by_key=backward_state_by_key,
                    visited_states=len(forward_visited | backward_visited),
                )

            meeting_key = self._expand_backward_layer(
                queue=backward_queue,
                visited=backward_visited,
                other_visited=forward_visited,
                parent=backward_parent,
                state_by_key=backward_state_by_key,
            )
            if meeting_key is not None:
                return self._build_result(
                    meeting_key=meeting_key,
                    start_key=start_key,
                    backward_root_keys=backward_root_keys,
                    forward_parent=forward_parent,
                    backward_parent=backward_parent,
                    forward_state_by_key=forward_state_by_key,
                    backward_state_by_key=backward_state_by_key,
                    visited_states=len(forward_visited | backward_visited),
                )

        return None

    def _expand_forward_layer(
        self,
        *,
        queue: deque[GameState],
        visited: set[tuple[object, ...]],
        other_visited: set[tuple[object, ...]],
        parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        state_by_key: dict[tuple[object, ...], GameState],
    ) -> tuple[object, ...] | None:
        for _ in range(len(queue)):
            current = queue.popleft()
            current_key = self.state_key(current)

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

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

                if next_key in other_visited:
                    return next_key

                queue.append(next_state)

        return None

    def _expand_backward_layer(
        self,
        *,
        queue: deque[GameState],
        visited: set[tuple[object, ...]],
        other_visited: set[tuple[object, ...]],
        parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        state_by_key: dict[tuple[object, ...], GameState],
    ) -> tuple[object, ...] | None:
        for _ in range(len(queue)):
            current = queue.popleft()
            current_key = self.state_key(current)

            for move, predecessor in self.iter_predecessors(current):
                predecessor_key = self.state_key(predecessor)
                if predecessor_key in visited:
                    continue

                visited.add(predecessor_key)
                # Stored move is the *forward* move predecessor -> current.
                parent[predecessor_key] = (current_key, move)
                state_by_key[predecessor_key] = predecessor

                if predecessor_key in other_visited:
                    return predecessor_key

                queue.append(predecessor)

        return None

    def _infer_forward_move(self, before: GameState, after: GameState) -> Move:
        """
        Because states are normalised, reverse-generation may reorder tubes.
        Rather than trying to hand-translate indices through that reorder, just
        infer the forward move directly by checking legal successors.
        """
        after_key = self.state_key(after)

        for move, next_state in self.iter_neighbours(before):
            if self.state_key(next_state) == after_key:
                return move

        raise ValueError("Could not infer a legal forward move between states.")

    def _build_result(
        self,
        *,
        meeting_key: tuple[object, ...],
        start_key: tuple[object, ...],
        backward_root_keys: set[tuple[object, ...]],
        forward_parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        backward_parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]],
        forward_state_by_key: dict[tuple[object, ...], GameState],
        backward_state_by_key: dict[tuple[object, ...], GameState],
        visited_states: int,
    ) -> SearchResult:
        forward_steps = self._reconstruct_forward_steps(
            parent=forward_parent,
            state_by_key=forward_state_by_key,
            start_key=start_key,
            goal_key=meeting_key,
        )
        backward_steps = self._reconstruct_backward_steps(
            parent=backward_parent,
            state_by_key=backward_state_by_key,
            start_key=meeting_key,
            goal_keys=backward_root_keys,
        )

        return SearchResult(
            steps=forward_steps + backward_steps,
            visited_states=visited_states,
        )

    @staticmethod
    def _reconstruct_forward_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)

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

        while current_key not in goal_keys:
            next_key, move = parent[current_key]
            path.append(SolutionStep(move=move, state=state_by_key[next_key]))
            current_key = next_key

        return tuple(path)
