"""Agent interfaces and wrappers for benchmarking.

This module defines the BenchmarkAgent protocol that all agents must
implement, plus wrapper classes for built-in agents (Random, Greedy, MCTS)
and neural network agents.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

import numpy as np
import torch

if TYPE_CHECKING:
    from numpy.typing import NDArray


@runtime_checkable
class BenchmarkAgent(Protocol):
    """Protocol for agents that can be benchmarked.

    All agents must implement:
    - select_action(): Choose action given observation and mask
    - reset(): Reset internal state for new game
    - name: Human-readable name for display

    Example:
        class MyAgent:
            @property
            def name(self) -> str:
                return "MyAgent"

            def select_action(self, obs: np.ndarray, mask: np.ndarray) -> int:
                # Your logic here
                valid_actions = np.where(mask > 0.5)[0]
                return int(np.random.choice(valid_actions))

            def reset(self) -> None:
                pass
    """

    @property
    def name(self) -> str:
        """Human-readable agent name."""
        ...

    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        """Select an action given observation and legal action mask.

        Args:
            observation: State tensor of shape (328,)
            action_mask: Binary mask of shape (422,), 1.0 = legal

        Returns:
            Action index in [0, 421]
        """
        ...

    def reset(self) -> None:
        """Reset agent state for a new game."""
        ...


class BaseAgent(ABC):
    """Abstract base class for benchmark agents."""

    def __init__(self, agent_name: str):
        self._name = agent_name

    @property
    def name(self) -> str:
        return self._name

    @abstractmethod
    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        pass

    def reset(self) -> None:  # noqa: B027
        """Default reset does nothing."""


class RandomAgent(BaseAgent):
    """Agent that selects random legal actions.

    This is the weakest baseline - any trained agent should beat it
    with >95% win rate.
    """

    def __init__(self, seed: int | None = None):
        super().__init__("Random")
        self.rng = np.random.default_rng(seed)

    def select_action(
        self,
        _observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        valid_actions = np.where(action_mask > 0.5)[0]
        if len(valid_actions) == 0:
            return 421  # EndTurn fallback
        return int(self.rng.choice(valid_actions))

    def reset(self) -> None:
        pass  # RNG continues across games


class GreedyAgent(BaseAgent):
    """Agent that uses the Rust GreedyBot via environment.

    This agent evaluates each action by simulation and picks the
    best according to a heuristic evaluation function.

    Win rates to expect:
    - vs Random: ~100%
    - vs Greedy: ~50% (mirror match)
    """

    def __init__(self) -> None:
        super().__init__("Greedy")
        # We'll use the Rust greedy bot through the game interface
        self._use_rust_bot = True

    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        # This is a placeholder - actual implementation uses Rust bot
        # through the benchmark runner which has access to the game state
        raise NotImplementedError(
            "GreedyAgent requires game state access. "
            "Use EssenceWarsBenchmark.evaluate() which handles this internally."
        )

    def reset(self) -> None:
        pass


class MCTSAgent(BaseAgent):
    """Agent that uses Monte Carlo Tree Search.

    MCTS with configurable simulation count provides a strong baseline.

    Win rates to expect (with 100 sims):
    - vs Random: ~100%
    - vs Greedy: ~60-70%
    - vs MCTS-50: ~55%
    """

    def __init__(self, simulations: int = 100):
        super().__init__(f"MCTS-{simulations}")
        self.simulations = simulations
        self._use_rust_bot = True

    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        # Placeholder - uses Rust MCTS through benchmark runner
        raise NotImplementedError(
            "MCTSAgent requires game state access. "
            "Use EssenceWarsBenchmark.evaluate() which handles this internally."
        )

    def reset(self) -> None:
        pass


class AlphaBetaAgent(BaseAgent):
    """Agent that uses Alpha-Beta minimax search.

    Alpha-Beta is a highly optimized search algorithm with:
    - Transposition table for position caching
    - Killer moves and history heuristic for move ordering
    - Aspiration windows for faster iterative deepening
    - Late Move Reduction (LMR) for deeper effective search

    This bot is 18-30x faster than MCTS at comparable strength.

    Win rates to expect (at depth 6):
    - vs Random: ~100%
    - vs Greedy: ~70-80%
    - vs MCTS-100: ~60-70%
    """

    def __init__(self, depth: int = 6):
        super().__init__(f"AlphaBeta-d{depth}")
        self.depth = depth
        self._use_rust_bot = True

    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        # Placeholder - uses Rust Alpha-Beta through benchmark runner
        raise NotImplementedError(
            "AlphaBetaAgent requires game state access. "
            "Use EssenceWarsBenchmark.evaluate() which handles this internally."
        )

    def reset(self) -> None:
        pass


class NeuralAgent(BaseAgent):
    """Agent that uses a neural network for action selection.

    Wraps any PyTorch model that takes observation and mask tensors
    and outputs action logits or probabilities.
    """

    def __init__(
        self,
        network: torch.nn.Module,
        name: str = "Neural",
        device: str | torch.device = "cpu",
        deterministic: bool = True,
        obs_normalizer: Any | None = None,
    ):
        """Initialize neural agent.

        Args:
            network: PyTorch model with forward(obs, mask) -> (logits, value)
            name: Display name for the agent
            device: Device to run inference on
            deterministic: If True, pick argmax; if False, sample from policy
            obs_normalizer: Optional RunningMeanStd for observation normalization
        """
        super().__init__(name)
        self.network = network
        self.device = torch.device(device)
        self.deterministic = deterministic
        self.obs_normalizer = obs_normalizer
        self.network.to(self.device)
        self.network.eval()

    def select_action(
        self,
        observation: NDArray[np.float32],
        action_mask: NDArray[np.float32],
    ) -> int:
        with torch.no_grad():
            if self.obs_normalizer is not None:
                observation = self.obs_normalizer.normalize(observation).astype(np.float32)
            obs = torch.from_numpy(observation).float().unsqueeze(0).to(self.device)
            mask = torch.from_numpy(action_mask).bool().unsqueeze(0).to(self.device)

            # Forward pass
            logits, _ = self.network(obs, mask)

            if self.deterministic:
                # Argmax over legal actions
                action = logits.argmax(dim=-1).item()
            else:
                # Sample from softmax distribution
                probs = torch.softmax(logits, dim=-1)
                action = torch.multinomial(probs, 1).item()

        return int(action)

    def reset(self) -> None:
        pass

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: str,
        name: str | None = None,
        device: str = "cpu",
        deterministic: bool = True,
    ) -> NeuralAgent:
        """Load a neural agent from a checkpoint file.

        Supports PPO, AlphaZero, and Behavioral Cloning checkpoints.

        Args:
            checkpoint_path: Path to .pt checkpoint file
            name: Display name (default: derived from filename)
            device: Device for inference
            deterministic: Whether to use argmax or sampling

        Returns:
            NeuralAgent instance

        Raises:
            FileNotFoundError: If checkpoint file doesn't exist
            ValueError: If checkpoint format is not recognized
        """
        from pathlib import Path

        # Check file exists with helpful error
        path = Path(checkpoint_path)
        if not path.exists():
            raise FileNotFoundError(
                f"Checkpoint file not found: {checkpoint_path}\n"
                "Make sure you've trained a model first. Example:\n"
                "  uv run python scripts/training/ppo.py --timesteps 10000"
            )

        if not path.suffix == ".pt":
            raise ValueError(
                f"Expected .pt checkpoint file, got: {path.suffix}\n"
                "Essence Wars uses PyTorch .pt checkpoint files."
            )

        try:
            checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        except Exception as e:
            raise ValueError(
                f"Failed to load checkpoint: {checkpoint_path}\n"
                f"Error: {e}\n"
                "This may indicate a corrupted or incompatible checkpoint file."
            ) from e

        if not isinstance(checkpoint, dict):
            raise ValueError(
                f"Invalid checkpoint format in: {checkpoint_path}\n"
                f"Expected a dict, got {type(checkpoint).__name__}.\n"
                "Essence Wars checkpoints should contain 'network_state_dict' or 'model_state_dict'."
            )

        # Extract config - handle both dict and dataclass configs
        config = checkpoint.get("config", {})
        config_dict: dict[str, Any]
        if hasattr(config, "__dict__"):
            # Dataclass config (e.g., PPOConfig) - vars() may return MappingProxyType
            config_dict = dict(vars(config))
        elif isinstance(config, dict):
            config_dict = dict(config)
        else:
            config_dict = {}

        # Determine network architecture from checkpoint
        if "model_state_dict" in checkpoint:
            # BC checkpoint
            state_dict = checkpoint["model_state_dict"]
            args = checkpoint.get("args", {})
            if hasattr(args, "__dict__"):
                args = dict(vars(args))
            config_dict.update(args)
        elif "network_state_dict" in checkpoint:
            # AlphaZero or PPO checkpoint
            state_dict = checkpoint["network_state_dict"]
        else:
            available_keys = list(checkpoint.keys())
            raise ValueError(
                f"Unknown checkpoint format: {checkpoint_path}\n"
                f"Expected 'network_state_dict' (PPO/AlphaZero) or 'model_state_dict' (BC).\n"
                f"Found keys: {available_keys}\n"
                "Make sure you're using a checkpoint from essence-wars training scripts."
            )

        # Get network parameters
        hidden_dim = config_dict.get("hidden_dim", 256)
        num_blocks = config_dict.get("num_blocks", 4)
        observation_mode = config_dict.get("observation_mode", "flat")
        embed_dim = config_dict.get("embed_dim", 64)

        from essence_wars._core import STATE_TENSOR_SIZE

        # Detect obs_dim from state dict weights (handles old 326-dim checkpoints)
        obs_dim = STATE_TENSOR_SIZE  # Default: current size (328)
        for key in ("input_proj.0.weight", "trunk.0.weight"):
            if key in state_dict:
                obs_dim = state_dict[key].shape[1]
                break

        # Determine network type from observation mode
        network: torch.nn.Module
        if observation_mode in ("embedded", "embedded_pretrained"):
            # PPO with embeddings
            from essence_wars.agents.embeddings import EmbeddedPPONetwork

            network = EmbeddedPPONetwork(
                embed_dim=embed_dim,
                hidden_dim=hidden_dim,
            )
        elif "num_blocks" in config_dict:
            # AlphaZero-style network (also used by BC)
            from essence_wars.agents.networks import AlphaZeroNetwork

            network = AlphaZeroNetwork(
                obs_dim=obs_dim,
                action_dim=422,
                hidden_dim=hidden_dim,
                num_blocks=num_blocks,
            )
        else:
            # Simple PPO network (flat observation)
            from essence_wars.agents.networks import EssenceWarsNetwork

            network = EssenceWarsNetwork(
                obs_dim=obs_dim,
                action_dim=422,
                hidden_dim=hidden_dim,
            )

        try:
            network.load_state_dict(state_dict)
        except RuntimeError as e:
            error_msg = str(e)
            if "Missing key" in error_msg and "input_norm" in error_msg:
                # Old checkpoint without input_norm — load with strict=False
                # LayerNorm defaults (weight=1, bias=0) won't match trained behavior,
                # so disable it by setting it to identity after loading
                network.load_state_dict(state_dict, strict=False)
                if hasattr(network, "input_norm"):
                    # Reset to pass-through: after LayerNorm normalizes, weight=1/bias=0
                    # is default, but we want true identity for old models
                    import torch.nn as nn
                    network.input_norm = nn.Identity()
            elif "size mismatch" in error_msg:
                raise ValueError(
                    f"Network architecture mismatch loading: {checkpoint_path}\n"
                    f"Error: {error_msg}\n"
                    "The checkpoint was trained with different network settings (hidden_dim, etc.).\n"
                    "Check that the checkpoint matches your expected configuration."
                ) from e
            elif "Missing key" in error_msg or "Unexpected key" in error_msg:
                raise ValueError(
                    f"State dict key mismatch loading: {checkpoint_path}\n"
                    f"Error: {error_msg}\n"
                    "The checkpoint may be from an incompatible model version."
                ) from e
            else:
                raise ValueError(
                    f"Failed to load state dict from: {checkpoint_path}\n"
                    f"Error: {error_msg}"
                ) from e

        # Restore observation normalizer if present
        obs_normalizer = None
        if "obs_normalizer" in checkpoint:
            from essence_wars.agents.utils.normalization import RunningMeanStd

            norm_data = checkpoint["obs_normalizer"]
            obs_normalizer = RunningMeanStd(shape=norm_data["mean"].shape)
            obs_normalizer.mean = norm_data["mean"]
            obs_normalizer.var = norm_data["var"]
            obs_normalizer.count = norm_data["count"]

        # Generate name from filename if not provided
        if name is None:
            name = Path(checkpoint_path).stem

        return cls(
            network,
            name=name,
            device=device,
            deterministic=deterministic,
            obs_normalizer=obs_normalizer,
        )
