import math
import random
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol


class GameEnv(Protocol):
    def current_player(self, state: Any) -> int: ...
    def legal_moves(self, state: Any) -> List[Any]: ...
    def next_state(self, state: Any, move: Any) -> Any: ...
    def is_terminal(self, state: Any) -> bool: ...
    def final_reward(self, state: Any, player: int) -> float: ...


@dataclass
class Node:
    state: Any
    player_to_move: int
    parent: Optional["Node"] = None
    move_from_parent: Optional[Any] = None

    children: Dict[Any, "Node"] = field(default_factory=dict)
    unexpanded_moves: List[Any] = field(default_factory=list)

    visits: int = 0
    value_sum: float = 0.0

    def q_value(self) -> float:
        return self.value_sum / self.visits if self.visits > 0 else 0.0

@dataclass
class MCTSContext:
    """
    Holds all data shared across MCTS phases during a single iteration.
    Used by the State pattern for Selection/Expansion/Simulation/Backprop.
    """
    env: GameEnv
    root: Node
    root_player: int

    selection_strategy: "SelectionStrategy"
    expansion_strategy: "ExpansionStrategy"
    simulation_strategy: "SimulationStrategy"
    backpropagation_strategy: "BackpropagationStrategy"

    current_node: Optional[Node] = None
    reward: float = 0.0







# Strategy Interfaces for MCTS Phases

class SelectionStrategy(ABC):
    """Strategy for selecting a node to expand during tree traversal."""

    @abstractmethod
    def select(self, node: Node, env: GameEnv) -> Node:
        """
        Select a node from the tree for expansion.

        Args:
            node: The current node to select from
            env: The game environment

        Returns:
            The selected node
        """
        pass


class ExpansionStrategy(ABC):
    """Strategy for expanding a node by adding a new child."""

    @abstractmethod
    def expand(self, node: Node, env: GameEnv) -> Optional[Node]:
        """
        Expand a node by adding a new child.

        Args:
            node: The node to expand
            env: The game environment

        Returns:
            The newly created child node, or None if node cannot be expanded
        """
        pass


class SimulationStrategy(ABC):
    """Strategy for simulating a game from a given node."""

    @abstractmethod
    def simulate(self, node: Node, env: GameEnv, root_player: int) -> float:
        """
        Simulate a game from the given node to a terminal state.

        Args:
            node: The node to simulate from
            env: The game environment
            root_player: The player from whose perspective to evaluate

        Returns:
            The reward for the root player
        """
        pass


class BackpropagationStrategy(ABC):
    """Strategy for backpropagating simulation results through the tree."""

    @abstractmethod
    def backpropagate(self, node: Node, reward: float) -> None:
        """
        Backpropagate the simulation result through the tree.

        Args:
            node: The node to start backpropagation from
            reward: The reward to backpropagate
        """
        pass


# Default Strategy Implementations

class UCBSelectionStrategy(SelectionStrategy):
    """UCB1-based selection strategy for tree traversal."""

    def __init__(self, exploration_c: float = 1.4):
        self.exploration_c = exploration_c

    def select(self, node: Node, env: GameEnv) -> Node:
        """Select child using UCB1 formula."""
        current = node

        while True:
            if env.is_terminal(current.state):
                return current

            # If node has unexpanded moves, return it for expansion
            if current.unexpanded_moves:
                return current

            # Otherwise, select best child using UCB
            current = self._ucb_select_child(current)

    def _ucb_select_child(self, node: Node) -> Node:
        """Select child with highest UCB value."""
        assert node.children, "Cannot select child of leaf without children"

        log_N = math.log(node.visits + 1)

        def ucb(child: Node) -> float:
            if child.visits == 0:
                return float("inf")
            exploit = child.q_value()
            explore = self.exploration_c * math.sqrt(log_N / child.visits)
            return exploit + explore

        _, best_child = max(
            node.children.items(),
            key=lambda item: ucb(item[1])
        )
        return best_child


class DefaultExpansionStrategy(ExpansionStrategy):
    """Default expansion strategy that expands one unexpanded move."""

    def expand(self, node: Node, env: GameEnv) -> Optional[Node]:
        """Expand by creating a child for one unexpanded move."""
        if not node.unexpanded_moves:
            return None

        move = node.unexpanded_moves.pop()
        next_state = env.next_state(node.state, move)
        child = Node(
            state=next_state,
            player_to_move=env.current_player(next_state),
            parent=node,
            move_from_parent=move,
            unexpanded_moves=env.legal_moves(next_state),
        )

        # ✅ Use a hashable key for the children dict
        if isinstance(move, dict):
            move_key = move.get("id", repr(move))
        else:
            move_key = move  # assume already hashable (e.g. tuple, int, str)

        node.children[move_key] = child
        return child


class RandomSimulationStrategy(SimulationStrategy):
    """Random rollout simulation strategy."""

    def simulate(self, node: Node, env: GameEnv, root_player: int) -> float:
        """Simulate game with random moves until terminal state."""
        state = node.state

        while not env.is_terminal(state):
            moves = env.legal_moves(state)
            if not moves:
                break
            move = random.choice(moves)
            state = env.next_state(state, move)

        return env.final_reward(state, root_player)


class DefaultBackpropagationStrategy(BackpropagationStrategy):
    """Default backpropagation strategy that updates visits and values."""

    def backpropagate(self, node: Node, reward: float) -> None:
        """Backpropagate reward up the tree."""
        current = node
        while current is not None:
            current.visits += 1
            current.value_sum += reward
            current = current.parent

class MCTSPhaseState(ABC):
    """
    State interface for a single phase of the MCTS iteration.
    Each concrete state performs its phase and returns the next state.
    """

    @abstractmethod
    def handle(self, context: MCTSContext) -> Optional["MCTSPhaseState"]:
        """
        Perform this phase's work on the context and return the next phase.
        Returning None indicates the iteration is finished.
        """
        pass


class SelectionPhaseState(MCTSPhaseState):
    """State for the Selection phase."""

    def handle(self, context: MCTSContext) -> Optional[MCTSPhaseState]:
        # Start at root each iteration, delegate to selection strategy
        node = context.selection_strategy.select(context.root, context.env)
        context.current_node = node
        return ExpansionPhaseState()


class ExpansionPhaseState(MCTSPhaseState):
    """State for the Expansion phase."""

    def handle(self, context: MCTSContext) -> Optional[MCTSPhaseState]:
        node = context.current_node
        if node is None:
            # Nothing to expand, go straight to simulation
            return SimulationPhaseState()

        expanded = context.expansion_strategy.expand(node, context.env)
        if expanded is not None:
            context.current_node = expanded

        return SimulationPhaseState()


class SimulationPhaseState(MCTSPhaseState):
    """State for the Simulation phase."""

    def handle(self, context: MCTSContext) -> Optional[MCTSPhaseState]:
        node = context.current_node
        if node is None:
            context.reward = 0.0
            return BackpropagationPhaseState()

        reward = context.simulation_strategy.simulate(
            node,
            context.env,
            context.root_player,
        )
        context.reward = reward
        return BackpropagationPhaseState()


class BackpropagationPhaseState(MCTSPhaseState):
    """State for the Backpropagation phase."""

    def handle(self, context: MCTSContext) -> Optional[MCTSPhaseState]:
        node = context.current_node
        if node is not None:
            context.backpropagation_strategy.backpropagate(node, context.reward)

        # End this MCTS iteration
        return None






# Main Search Algorithm Strategy

class SearchAlgorithmStrategy(ABC):
    """Strategy interface for tree search algorithms."""

    @abstractmethod
    def search(self, root_state: Any, iterations: int, root_player: Optional[int] = None) -> Any:
        """
        Search for the best move from the given state.

        Args:
            root_state: The state to search from
            iterations: Number of search iterations to perform
            root_player: The player to search for (if None, determined from state)

        Returns:
            The best move found
        """
        pass


class MCTSStrategy(SearchAlgorithmStrategy):
    """Monte Carlo Tree Search algorithm implementation using pluggable strategies."""

    def __init__(
        self,
        env: GameEnv,
        selection_strategy: Optional[SelectionStrategy] = None,
        expansion_strategy: Optional[ExpansionStrategy] = None,
        simulation_strategy: Optional[SimulationStrategy] = None,
        backpropagation_strategy: Optional[BackpropagationStrategy] = None,
        exploration_c: float = 1.4,
    ):
        self.env = env
        self.selection_strategy = selection_strategy or UCBSelectionStrategy(exploration_c)
        self.expansion_strategy = expansion_strategy or DefaultExpansionStrategy()
        self.simulation_strategy = simulation_strategy or RandomSimulationStrategy()
        self.backpropagation_strategy = backpropagation_strategy or DefaultBackpropagationStrategy()

    def search(self, root_state: Any, iterations: int, root_player: Optional[int] = None) -> Any:
        """Execute MCTS search using the configured strategies."""
        if root_player is None:
            root_player = self.env.current_player(root_state)

        root = Node(
            state=root_state,
            player_to_move=self.env.current_player(root_state),
            unexpanded_moves=self.env.legal_moves(root_state),
        )

        # Shared context for all iterations
        context = MCTSContext(
            env=self.env,
            root=root,
            root_player=root_player,
            selection_strategy=self.selection_strategy,
            expansion_strategy=self.expansion_strategy,
            simulation_strategy=self.simulation_strategy,
            backpropagation_strategy=self.backpropagation_strategy,
        )

        for _ in range(iterations):
            # Start each iteration in the Selection phase
            state: Optional[MCTSPhaseState] = SelectionPhaseState()
            while state is not None:
                state = state.handle(context)

        if not root.children:
            return None

        # Pick child with most visits
        _, best_child = max(
            root.children.items(),
            key=lambda item: item[1].visits,
        )

        # ✅ Return the original move dict stored in the child
        return best_child.move_from_parent


class MCTS:
    """
    Monte Carlo Tree Search - Backward compatibility wrapper.

    This class wraps MCTSStrategy to maintain backward compatibility
    with existing code while using the new strategy pattern internally.

    For new code, consider using MCTSStrategy directly for more flexibility.
    """

    def __init__(
        self,
        env: GameEnv,
        exploration_c: float = 1.4,
        selection_strategy: Optional[SelectionStrategy] = None,
        expansion_strategy: Optional[ExpansionStrategy] = None,
        simulation_strategy: Optional[SimulationStrategy] = None,
        backpropagation_strategy: Optional[BackpropagationStrategy] = None,
    ):
        """
        Initialize MCTS with optional custom strategies.

        Args:
            env: The game environment
            exploration_c: Exploration constant for UCB (default: 1.4)
            selection_strategy: Custom selection strategy (default: UCBSelectionStrategy)
            expansion_strategy: Custom expansion strategy (default: DefaultExpansionStrategy)
            simulation_strategy: Custom simulation strategy (default: RandomSimulationStrategy)
            backpropagation_strategy: Custom backpropagation strategy (default: DefaultBackpropagationStrategy)
        """
        self.env = env
        self.exploration_c = exploration_c
        self._strategy = MCTSStrategy(
            env=env,
            selection_strategy=selection_strategy,
            expansion_strategy=expansion_strategy,
            simulation_strategy=simulation_strategy,
            backpropagation_strategy=backpropagation_strategy,
            exploration_c=exploration_c,
        )

    def search(self, root_state: Any, iterations: int, root_player: Optional[int] = None) -> Any:
        """
        Search for the best move from the given state.

        Args:
            root_state: The state to search from
            iterations: Number of search iterations to perform
            root_player: The player to search for (if None, determined from state)

        Returns:
            The best move found
        """
        return self._strategy.search(root_state, iterations, root_player)
