from __future__ import annotations
import math
import random
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Callable

from .client import GameClient

State = Dict[str, Any]
Move = Dict[str, Any]
ApplyMoveFn = Callable[[State, Move], State]


@dataclass
class _Node:
    state: State
    player_to_move: int
    parent: Optional["_Node"] = None
    move_from_parent: Optional[Move] = None

    children: Dict[str, "_Node"] = field(default_factory=dict)
    unexpanded_moves: List[Move] = 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 _GenericGameEnv:
    def __init__(self, apply_move_fn: ApplyMoveFn):
        self.apply_move_fn = apply_move_fn

    def current_player(self, state: State) -> int:
        return state["turn_index"]

    def legal_moves(self, state: State) -> List[Move]:
        moves = state.get("legal_moves")
        if moves is None:
            raise ValueError("state['legal_moves'] is required for MCTS.")
        return moves

    def next_state(self, state: State, move: Move) -> State:
        return self.apply_move_fn(state, move)

    def is_terminal(self, state: State) -> bool:
        return bool(state.get("is_terminal", False))

    def final_reward(self, state: State, player_index: int) -> float:
        result = state.get("result")
        if result is None:
            return 0.0

        players = state["players"]
        player_id = players[player_index]["id"]
        return float(result.get(player_id, 0.0))


class _MCTSCore:
    def __init__(self, env: _GenericGameEnv, exploration_c: float = 1.4):
        self.env = env
        self.exploration_c = exploration_c

    def search(self, root_state: State, iterations: int) -> Optional[Move]:
        root_player = self.env.current_player(root_state)
        root = _Node(
            state=root_state,
            player_to_move=root_player,
            unexpanded_moves=list(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_child = max(
            root.children.items(),
            key=lambda item: item[1].visits,
        )
        return best_child.move_from_parent

    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=list(self.env.legal_moves(next_state)),
                )
                move_key = move.get("id", str(move))
                node.children[move_key] = child
                return child

            node = self._ucb_select_child(node)

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


class AIGameClient(GameClient):
    """
    Unified client:
      - All GameClient methods (start_match, log_move, end_match, ...)
      - Plus MCTS AI:
          send_state(match_id, state)
          best_move(match_id, iterations=1000)
    """

    def __init__(
        self,
        game_id: str,
        api_key: str,
        apply_move_fn: ApplyMoveFn,
        base_url: str = "http://localhost:8000",
        exploration_c: float = 1.4,
    ):
        super().__init__(game_id=game_id, api_key=api_key, base_url=base_url)
        self._env = _GenericGameEnv(apply_move_fn)
        self._mcts = _MCTSCore(self._env, exploration_c=exploration_c)
        self._states: Dict[str, State] = {}  # match_id -> last state

    def send_state(self, match_id: str, state: State) -> None:
        """
        Store the latest game state for a match.
        The state must be in the generic format produced by build_generic_state.
        """
        if state["game_id"] != self.game_id:
            raise ValueError(
                f"State game_id={state['game_id']} does not match client.game_id={self.game_id}"
            )
        self._states[match_id] = state

    def best_move(self, match_id: str, iterations: int = 1000) -> Optional[Move]:
        """
        Run MCTS from the last state sent for this match and return
        one of the moves from state['legal_moves'].
        """
        state = self._states.get(match_id)
        if state is None:
            raise KeyError(f"No state stored for match_id={match_id}")

        if self._env.is_terminal(state):
            return None

        return self._mcts.search(state, iterations=iterations)
