from __future__ import annotations
from typing import Any, Dict, List, Optional, Callable, Protocol, TypeVar, Union
from .mcts import (
    MCTSStrategy,
    GameEnv,
    SelectionStrategy,
    ExpansionStrategy,
    SimulationStrategy,
    BackpropagationStrategy,
    RandomSimulationStrategy,
    get_simulation_strategy,
    get_selection_strategy,
    get_backpropagation_strategy,
    get_expansion_strategy,
)


from .client import GameClient
from .rabbitmq import RabbitMQ

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

class TurnBasedGame(Protocol):
    """
    Generic interface that any turn-based game must implement
    to work with our MCTS adapter.
    """

    def get_legal_actions(self) -> List[Move]:
        ...

    def is_game_over(self) -> bool:
        ...

    def game_result(self) -> int:
        """
        Typically 1 / 0 / -1 from the perspective of the player to move.
        The SDK itself doesn’t care; this is mainly for tutorial-style logic.
        """
        ...

    def move(self, action: Move) -> "TurnBasedGame":
        """
        Apply an action and return a NEW game state.
        """
        ...


class GenericGameEnv(GameEnv):
    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 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,
        selection_strategy: Optional[Union[str, SelectionStrategy]] = None,
        expansion_strategy: Optional[Union[str, ExpansionStrategy]] = None,
        simulation_strategy: Optional[Union[str, SimulationStrategy]] = None,
        backpropagation_strategy: Optional[Union[str, BackpropagationStrategy]] = None,
        queue_name: Optional[str] = None,
    ):
        try:
            message_bus = RabbitMQ()
            print("[AIGameClient] Connected to RabbitMQ.")
        except Exception as e:
            print(
                 f"[AIGameClient] WARNING: could not connect to RabbitMQ ({e}). "
                "Falling back to stdout logging."
            )
            message_bus = None

        super().__init__(game_id=game_id, api_key=api_key, base_url=base_url, message_bus=message_bus)
        self._queue_name = queue_name or "game_events"

        # Concrete GameEnv implementation for generic JSON states
        self._env: GameEnv = GenericGameEnv(apply_move_fn)

        # Resolve string-based strategies to instances
        resolved_selection: Optional[SelectionStrategy] = None
        if isinstance(selection_strategy, str):
            resolved_selection = get_selection_strategy(selection_strategy, exploration_c)
        else:
            resolved_selection = selection_strategy

        resolved_simulation: Optional[SimulationStrategy] = None
        if isinstance(simulation_strategy, str):
            resolved_simulation = get_simulation_strategy(simulation_strategy)
            # RAVE needs move tracking enabled
            if isinstance(resolved_simulation, RandomSimulationStrategy):
                if isinstance(resolved_selection, str) and resolved_selection.lower() == "rave":
                    resolved_simulation.track_moves = True
                elif hasattr(resolved_selection, '__class__') and 'RAVE' in resolved_selection.__class__.__name__:
                    resolved_simulation.track_moves = True
        else:
            resolved_simulation = simulation_strategy

        resolved_backprop: Optional[BackpropagationStrategy] = None
        if isinstance(backpropagation_strategy, str):
            resolved_backprop = get_backpropagation_strategy(backpropagation_strategy)
        else:
            resolved_backprop = backpropagation_strategy

        resolved_expansion: Optional[ExpansionStrategy] = None
        if isinstance(expansion_strategy, str):
            resolved_expansion = get_expansion_strategy(expansion_strategy)
        else:
            resolved_expansion = expansion_strategy

        # Plug in MCTSStrategy with optional custom phase strategies
        self._mcts = MCTSStrategy(
            env=self._env,
            selection_strategy=resolved_selection,
            expansion_strategy=resolved_expansion,
            simulation_strategy=resolved_simulation,
            backpropagation_strategy=resolved_backprop,
            exploration_c=exploration_c,
        )

        # match_id -> last state
        self._states: Dict[str, 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

        # Uses strategy-based MCTS under the hood
        return self._mcts.search(root_state=state, iterations=iterations)

GameT = TypeVar("GameT", bound=TurnBasedGame)


def make_apply_move_fn(
    state_to_game: Callable[[State], GameT],
    game_to_state: Callable[[GameT, State], State],
) -> ApplyMoveFn:
    """
    Build an ApplyMoveFn from two small adapters:

      - state_to_game(state) -> Game
      - game_to_state(game, previous_state) -> new State

    Works for ANY game that implements TurnBasedGame.
    """

    def apply_move(state: State, move: Move) -> State:
        # 1) Convert generic state dict into a concrete game object
        game = state_to_game(state)

        # 2) Let the game apply the move using its own rules
        next_game = game.move(move)

        # 3) Convert the updated game back into generic JSON-style state
        new_state = game_to_state(next_game, state)
        return new_state

    return apply_move
