import math
import random
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


class MCTS:
    def __init__(self, env: GameEnv, exploration_c: float = 1.4):
        self.env = env
        self.exploration_c = exploration_c

    def search(self, root_state: Any, iterations: int, root_player: Optional[int] = None) -> Any:
        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),
        )

        for _ in range(iterations):
            node = self._select(root)
            reward = self._simulate(node, root_player)
            self._backpropagate(node, reward)

        if not root.children:
            return None

        best_move, best_child = max(
            root.children.items(),
            key=lambda item: item[1].visits
        )
        return best_move

    def _select(self, node: Node) -> Node:
        while True:
            if self.env.is_terminal(node.state):
                return node

            if node.unexpanded_moves:
                move = node.unexpanded_moves.pop()
                next_state = self.env.next_state(node.state, move)
                child = Node(
                    state=next_state,
                    player_to_move=self.env.current_player(next_state),
                    parent=node,
                    move_from_parent=move,
                    unexpanded_moves=self.env.legal_moves(next_state),
                )
                node.children[move] = child
                return child

            node = self._ucb_select_child(node)

    def _ucb_select_child(self, node: Node) -> Node:
        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

    def _simulate(self, node: Node, root_player: int) -> float:
        state = node.state

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

        return self.env.final_reward(state, root_player)

    def _backpropagate(self, node: Node, reward: float) -> None:
        while node is not None:
            node.visits += 1
            node.value_sum += reward
            node = node.parent
