"""
A* solver for Dr Eureka with an admissible heuristic.

The heuristic counts how many tubes in the current state cannot be matched
to an identical tube in the goal state (by content and orientation).
Each such tube requires at least one move (pour or flip), so the heuristic
never overestimates the remaining cost.
"""

import contextlib
import heapq
from collections import Counter
from collections.abc import Iterable

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


def _tube_key(tube: Tube) -> tuple[bool, tuple[str, ...]]:
    """Hashable key for a tube (orientation + ball sequence)."""
    return (tube.base_up, tuple(b.value for b in tube.balls))


def heuristic(current: GameState, goal: GameState) -> int:
    """
    Admissible heuristic: minimum number of moves needed from current to goal.

    Counts how many tubes in the current state have no identical match in the
    goal state (same balls, same base_up). Each such tube needs at least one
    move (pour or flip), so this is a lower bound on remaining moves.
    """
    goal_counts: Counter[tuple[bool, tuple[str, ...]]] = Counter()
    for tube in goal.canonical():
        goal_counts[_tube_key(tube)] += 1

    h = 0
    for tube in current.canonical():
        key = _tube_key(tube)
        if goal_counts[key] > 0:
            goal_counts[key] -= 1
        else:
            h += 1
    return h


class AStarCursorSolver(DrEurekaSolver):
    """
    A* solver for Dr Eureka.

    Uses a priority queue ordered by f = g + h, where g is moves from start
    and h is the number of tubes that cannot be matched to the goal.
    """

    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)

        # Priority queue: (f, tiebreaker, state_key). We only store keys to avoid
        # comparing Tube/GameState in the heap (they may not be orderable).
        tiebreaker = 0
        g_start = 0
        h_start = heuristic(start, goal)
        f_start = g_start + h_start

        open_heap: list[tuple[int, int, tuple[object, ...]]] = []
        heapq.heappush(open_heap, (f_start, tiebreaker, start_key))
        tiebreaker += 1

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

        while open_heap:
            _, _, current_key = heapq.heappop(open_heap)
            current = state_by_key[current_key]
            g = g_score[current_key]

            if current_key in visited:
                continue
            visited.add(current_key)

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

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

                tentative_g = g + 1
                if tentative_g < g_score.get(next_key, float("inf")):
                    g_score[next_key] = tentative_g
                    parent[next_key] = (current_key, move)
                    state_by_key[next_key] = next_state
                    h = heuristic(next_state, goal)
                    f_next = tentative_g + h
                    heapq.heappush(open_heap, (f_next, tiebreaker, next_key))
                    tiebreaker += 1

        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)


class AStarGippitySolver(DrEurekaSolver):
    """
    A* solver for Dr Eureka.

    The heuristic below is intentionally simple but admissible.

    Reasoning:
    - We compare the current state and goal state using their canonical
      representations, so tube positions do not matter.
    - Count how many current tubes cannot be matched exactly to a distinct
      goal tube.
    - One legal move can change at most two tubes:
        * a pour changes source and destination
        * a flip changes one tube
    - Therefore a single move can reduce the number of mismatched tubes by
      at most 2.
    - So ceil(mismatched_tubes / 2) is a lower bound on the number of
      remaining moves.

    This is a weak heuristic, but it is cheap and safe.
    """

    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 heuristic(self, state: GameState, goal: GameState) -> int:
        current_tubes = list(state.canonical())
        goal_tubes = list(goal.canonical())

        used_goal_indices: set[int] = set()
        exact_matches = 0

        for tube in current_tubes:
            for goal_index, goal_tube in enumerate(goal_tubes):
                if goal_index in used_goal_indices:
                    continue
                if tube == goal_tube:
                    used_goal_indices.add(goal_index)
                    exact_matches += 1
                    break

        mismatched_tubes = len(current_tubes) - exact_matches
        return (mismatched_tubes + 1) // 2

    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)

        open_heap: list[tuple[int, int, int, GameState]] = []
        g_score: dict[tuple[object, ...], int] = {start_key: 0}
        parent: dict[tuple[object, ...], tuple[tuple[object, ...], Move]] = {}
        state_by_key: dict[tuple[object, ...], GameState] = {start_key: start}
        closed: set[tuple[object, ...]] = set()

        tie_breaker = 0
        start_h = self.heuristic(start, goal)
        heapq.heappush(open_heap, (start_h, 0, tie_breaker, start))

        while open_heap:
            _, current_g, _, current = heapq.heappop(open_heap)
            current_key = self.state_key(current)

            if current_key in closed:
                continue

            best_known_g = g_score.get(current_key)
            if best_known_g is None or current_g != best_known_g:
                continue

            if current_key == goal_key:
                steps = self._reconstruct_steps(parent, state_by_key, start_key, current_key)
                return SearchResult(
                    steps=steps,
                    visited_states=len(closed) + 1,
                )

            closed.add(current_key)

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

                tentative_g = current_g + 1
                known_g = g_score.get(next_key)
                if known_g is not None and tentative_g >= known_g:
                    continue

                g_score[next_key] = tentative_g
                parent[next_key] = (current_key, move)
                state_by_key[next_key] = next_state

                tie_breaker += 1
                next_h = self.heuristic(next_state, goal)
                heapq.heappush(
                    open_heap,
                    (tentative_g + next_h, tentative_g, tie_breaker, 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)
