from __future__ import annotations
from abc import ABC, abstractmethod
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 .dynamic_mcts_agent import DynamicMCTSAgent
from .utils import build_generic_state

from .client import GameClient
from .rabbitmq import RabbitMQ

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

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

    Subclasses must implement:
    - get_legal_actions()
    - is_game_over()
    - game_result()
    - move(action)

    Subclasses automatically get:
    - ai_vs_ai_difficulty_selection() for AI vs AI games with different difficulties
    """

    @abstractmethod
    def get_legal_actions(self) -> List[Move]:
        """Return list of legal moves in the current state."""
        ...

    @abstractmethod
    def is_game_over(self) -> bool:
        """Return True if the game has ended."""
        ...

    @abstractmethod
    def game_result(self) -> int:
        """
        Return game result from the perspective of the player to move.
        Typically 1 (win), 0 (draw), or -1 (loss).
        """
        ...

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

    def ai_vs_ai_difficulty_selection(
        self,
        difficulty1: str,
        difficulty2: str,
        state_to_game_fn: Callable[[State], "TurnBasedGame"],
        game_to_state_fn: Callable[["TurnBasedGame", State], State],
        game_id: str = "game",
        verbose: bool = True
    ) -> Dict[str, Any]:
        """
        Play a complete AI vs AI game with different difficulty levels.

        This method allows you to run a full game where two AI agents with different
        difficulty settings play against each other using a single agent instance.

        Args:
            difficulty1: Difficulty for player 1 ("easy", "medium", or "hard")
            difficulty2: Difficulty for player 2 ("easy", "medium", or "hard")
            state_to_game_fn: Function to convert state dict to game object
            game_to_state_fn: Function to convert game object to state dict
            game_id: ID of the game (default: "game")
            verbose: Whether to print game progress (default: True)

        Returns:
            Dictionary with game results containing:
            - winner: Winning player symbol or None for draw
            - move_count: Total number of moves played
            - move_history: List of move records with player info and positions
            - final_game: Final game state object

        Example:
            # Create initial game
            game = TicTacToeGame(empty_board, players, "X")

            # Play AI vs AI with different difficulties
            result = game.ai_vs_ai_difficulty_selection(
                difficulty1="easy",
                difficulty2="hard",
                state_to_game_fn=state_to_game_tictactoe,
                game_to_state_fn=game_to_state_tictactoe,
                game_id="tictactoe"
            )

            print(f"Winner: {result['winner']}")
            print(f"Moves: {result['move_count']}")
        """
        # Create the apply_move function
        def apply_move_fn(state: State, move: Move) -> State:
            game = state_to_game_fn(state)
            next_game = game.move(move)
            return game_to_state_fn(next_game, state)

        # Create environment and agent
        env = GenericGameEnv(apply_move_fn)
        agent = DynamicMCTSAgent(env)

        # Assume the game has these attributes (duck typing)
        # If your game doesn't have these, you'll need to adapt
        players = getattr(self, 'players', [
            {"id": "P1", "symbol": "Player1"},
            {"id": "P2", "symbol": "Player2"}
        ])

        # Assign difficulties to players
        players_with_difficulty = [
            {**players[0], "difficulty": difficulty1},
            {**players[1], "difficulty": difficulty2}
        ]

        current_game = self
        move_history = []
        move_count = 0
        current_player_idx = 0

        if verbose:
            print("=" * 70)
            print(f"AI vs AI Game: {difficulty1.upper()} vs {difficulty2.upper()}")
            print("=" * 70)
            print()

        # Game loop
        while not current_game.is_game_over():
            current_player = players_with_difficulty[current_player_idx]
            current_difficulty = current_player["difficulty"]

            # Get legal moves
            legal_moves = current_game.get_legal_actions()
            if not legal_moves:
                break

            # Build current state
            current_player_symbol = getattr(current_game, 'current_player_symbol',
                                           players_with_difficulty[current_player_idx]["symbol"])

            state = game_to_state_fn(current_game, {
                "game_id": game_id,
                "players": players_with_difficulty,
                "turn_index": current_player_idx,
                "move_count": move_count,
                "is_terminal": False,
                "legal_moves": legal_moves
            })

            # Select move at current difficulty
            if verbose:
                print(f"Move {move_count + 1}: {current_player['id']} at '{current_difficulty}' difficulty")

            move = agent.select_move_at_difficulty(state, current_difficulty)

            if not move:
                if verbose:
                    print("  No valid move found!")
                break

            if verbose:
                pos = move.get("position", move)
                print(f"  Selected: {pos}")

            # Record move
            move_history.append({
                "player": current_player["id"],
                "move": move,
                "difficulty": current_difficulty
            })

            # Apply move
            current_game = current_game.move(move)
            move_count += 1

            # Switch players
            current_player_idx = 1 - current_player_idx

            if verbose:
                print()

        # Determine winner
        winner = None
        if current_game.is_game_over():
            result = current_game.game_result()
            if result == 1:
                # Current player to move would win, but game is over, so previous player won
                winner_idx = 1 - current_player_idx
                winner = players_with_difficulty[winner_idx].get("symbol")
            elif result == -1:
                # Current player to move would lose, so they lost (shouldn't happen)
                winner_idx = current_player_idx
                winner = players_with_difficulty[winner_idx].get("symbol")
            # result == 0 means draw, winner stays None

        if verbose:
            print("=" * 70)
            print("Game Over!")
            print("=" * 70)
            if winner:
                winner_player = next(p for p in players_with_difficulty
                                   if p.get("symbol") == winner)
                print(f"Winner: {winner_player['id']} ({winner}) at '{winner_player['difficulty']}' difficulty")
            else:
                print("Result: Draw")
            print(f"Total moves: {move_count}")
            print()

        return {
            "winner": winner,
            "move_count": move_count,
            "move_history": move_history,
            "final_game": current_game
        }


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,
        move_evaluator: Optional[Callable] = 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, move_evaluator=move_evaluator)
            # 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
