"""Unified evaluator for Essence Wars agents.

Supports both neural network agents (from registry) and algorithmic bots.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np

if TYPE_CHECKING:
    from essence_wars.benchmark.agents import BenchmarkAgent



@dataclass
class AgentSpec:
    """Specification for an agent to evaluate.

    Attributes:
        name: Display name for the agent
        agent_type: Type of agent (neural, greedy, mcts, alphabeta, random)
        config: Configuration (e.g., simulations for MCTS, depth for AB)
        path: Path to model checkpoint (for neural agents)
        registry_name: Name in model registry (for neural agents)
    """

    name: str
    agent_type: str
    config: dict[str, Any] = field(default_factory=dict)
    path: str | None = None
    registry_name: str | None = None


def parse_agent_spec(spec: str) -> AgentSpec:
    """Parse an agent specification string.

    Supported formats:
        random          -> RandomAgent
        greedy          -> GreedyBot
        mcts-100        -> MCTS with 100 simulations
        alphabeta-6     -> Alpha-Beta with depth 6
        ppo-generalist  -> Neural agent from registry
        path:some/model.pt -> Neural agent from explicit path

    Args:
        spec: Agent specification string

    Returns:
        AgentSpec describing the agent
    """
    spec = spec.strip().lower()

    # Explicit path
    if spec.startswith("path:"):
        path = spec[5:]
        name = Path(path).stem
        return AgentSpec(
            name=name,
            agent_type="neural",
            path=path,
        )

    # Random
    if spec == "random":
        return AgentSpec(name="Random", agent_type="random")

    # Greedy
    if spec == "greedy":
        return AgentSpec(name="Greedy", agent_type="greedy")

    # MCTS with simulations
    if spec.startswith("mcts"):
        sims_str = spec.replace("mcts-", "").replace("mcts", "")
        sims = int(sims_str) if sims_str else 100
        return AgentSpec(
            name=f"MCTS-{sims}",
            agent_type="mcts",
            config={"simulations": sims},
        )

    # Alpha-Beta with depth
    if spec.startswith("alphabeta") or spec.startswith("ab"):
        depth_str = spec.replace("alphabeta-", "").replace("alphabeta", "")
        depth_str = depth_str.replace("ab-", "").replace("ab", "")
        depth_str = depth_str.replace("d", "")
        depth = int(depth_str) if depth_str else 6
        return AgentSpec(
            name=f"AlphaBeta-d{depth}",
            agent_type="alphabeta",
            config={"depth": depth},
        )

    # Assume it's a registry name (e.g., ppo-generalist)
    return AgentSpec(
        name=spec,
        agent_type="neural",
        registry_name=spec,
    )


@dataclass
class MatchResult:
    """Result of a single matchup between two agents."""

    agent1: str
    agent2: str
    games: int
    agent1_wins: int
    agent2_wins: int
    draws: int

    @property
    def agent1_win_rate(self) -> float:
        return self.agent1_wins / self.games if self.games > 0 else 0.0

    @property
    def agent2_win_rate(self) -> float:
        return self.agent2_wins / self.games if self.games > 0 else 0.0


@dataclass
class EvaluationResults:
    """Results from a unified evaluation run."""

    agents: list[str]
    matchups: list[MatchResult]
    elo_ratings: dict[str, float]
    total_games: int
    elapsed_seconds: float

    def win_rate_matrix(self) -> dict[str, dict[str, float]]:
        """Get win rate matrix as nested dict."""
        matrix: dict[str, dict[str, float]] = {a: {} for a in self.agents}
        for match in self.matchups:
            matrix[match.agent1][match.agent2] = match.agent1_win_rate
            matrix[match.agent2][match.agent1] = match.agent2_win_rate
        return matrix

    def summary(self) -> str:
        """Generate human-readable summary."""
        lines = [
            "Evaluation Results",
            "=" * 50,
            f"Agents: {len(self.agents)}",
            f"Total games: {self.total_games}",
            f"Time: {self.elapsed_seconds:.1f}s",
            "",
            "ELO Ratings:",
        ]

        # Sort by ELO descending
        sorted_agents = sorted(
            self.agents, key=lambda a: self.elo_ratings.get(a, 1500), reverse=True
        )
        for agent in sorted_agents:
            elo = self.elo_ratings.get(agent, 1500)
            lines.append(f"  {agent:<25} {elo:.0f}")

        return "\n".join(lines)


class UnifiedEvaluator:
    """Unified evaluator for neural and algorithmic agents.

    Example:
        evaluator = UnifiedEvaluator()
        results = evaluator.evaluate(["ppo-generalist", "greedy", "mcts-100"])
        print(results.summary())
    """

    def __init__(
        self,
        games_per_matchup: int = 50,
        verbose: bool = True,
        update_elo: bool = False,
    ):
        """Initialize evaluator.

        Args:
            games_per_matchup: Number of games per agent pair
            verbose: Print progress during evaluation
            update_elo: Unused, kept for API compatibility
        """
        self.games_per_matchup = games_per_matchup
        self.verbose = verbose
        self.update_elo = update_elo

        # Simple ELO ratings tracked in-memory for this evaluation run
        self._elo_ratings: dict[str, float] = {}

        # Cache for loaded agents
        self._agent_cache: dict[str, BenchmarkAgent] = {}

    def _update_elo(self, name1: str, name2: str, result: int, k: float = 32.0) -> None:
        """Update ELO ratings for a single game result."""
        r1 = self._elo_ratings.setdefault(name1, 1500.0)
        r2 = self._elo_ratings.setdefault(name2, 1500.0)
        e1 = 1.0 / (1.0 + 10.0 ** ((r2 - r1) / 400.0))
        if result == 0:
            s1 = 1.0
        elif result == 1:
            s1 = 0.0
        else:
            s1 = 0.5
        self._elo_ratings[name1] = r1 + k * (s1 - e1)
        self._elo_ratings[name2] = r2 + k * ((1.0 - s1) - (1.0 - e1))

    def _log(self, msg: str) -> None:
        if self.verbose:
            print(msg)

    def check_availability(
        self, agent_specs: list[str]
    ) -> tuple[list[AgentSpec], list[AgentSpec]]:
        """Check which agents are available.

        Args:
            agent_specs: List of agent specification strings

        Returns:
            Tuple of (available, missing) AgentSpec lists
        """
        available = []
        missing = []

        for spec_str in agent_specs:
            spec = parse_agent_spec(spec_str)

            if spec.agent_type == "neural":
                if spec.path and Path(spec.path).exists():
                    available.append(spec)
                else:
                    missing.append(spec)
            else:
                # Algorithmic bots are always available
                available.append(spec)

        return available, missing

    def _load_agent(self, spec: AgentSpec) -> BenchmarkAgent:
        """Load an agent from its specification."""
        from essence_wars.benchmark.agents import (
            AlphaBetaAgent,
            GreedyAgent,
            MCTSAgent,
            NeuralAgent,
            RandomAgent,
        )

        cache_key = f"{spec.agent_type}:{spec.name}:{spec.path}"
        if cache_key in self._agent_cache:
            return self._agent_cache[cache_key]

        agent: BenchmarkAgent

        if spec.agent_type == "random":
            agent = RandomAgent()
        elif spec.agent_type == "greedy":
            agent = GreedyAgent()
        elif spec.agent_type == "mcts":
            agent = MCTSAgent(simulations=spec.config.get("simulations", 100))
        elif spec.agent_type == "alphabeta":
            agent = AlphaBetaAgent(depth=spec.config.get("depth", 6))
        elif spec.agent_type == "neural":
            if not spec.path:
                raise ValueError(f"Neural agent {spec.name} has no path")
            agent = NeuralAgent.from_checkpoint(spec.path, name=spec.name)
        else:
            raise ValueError(f"Unknown agent type: {spec.agent_type}")

        self._agent_cache[cache_key] = agent
        return agent

    def _play_game(
        self,
        agent1: BenchmarkAgent,
        agent2: BenchmarkAgent,
        spec1: AgentSpec,
        spec2: AgentSpec,
        seed: int,
    ) -> int:
        """Play a single game between two agents.

        Returns:
            0 if agent1 won, 1 if agent2 won, -1 for draw
        """
        from essence_wars._core import PyGame

        # Select random decks
        decks = PyGame.list_decks()
        deck1 = np.random.choice(decks)
        deck2 = np.random.choice(decks)

        game = PyGame(deck1=deck1, deck2=deck2, game_mode="attrition")
        game.reset(seed=seed)

        agent1.reset()
        agent2.reset()

        turns = 0
        while not game.is_done() and turns < 500:
            current_player = game.current_player()
            obs = np.array(game.observe(), dtype=np.float32)
            mask = np.array(game.action_mask(), dtype=np.float32)

            if current_player == 0:
                # Agent 1's turn
                action = self._get_action(agent1, spec1, game, obs, mask)
            else:
                # Agent 2's turn
                action = self._get_action(agent2, spec2, game, obs, mask)

            game.step(action)
            turns += 1

        # Determine winner
        reward = game.get_reward(0)
        if reward > 0:
            return 0  # Agent 1 won
        elif reward < 0:
            return 1  # Agent 2 won
        else:
            return -1  # Draw

    def _get_action(
        self,
        agent: BenchmarkAgent,
        spec: AgentSpec,
        game: Any,
        obs: np.ndarray[Any, Any],
        mask: np.ndarray[Any, Any],
    ) -> int:
        """Get action from an agent, handling Rust bot special cases."""
        # For Rust bots, use the game's built-in methods
        if spec.agent_type == "random":
            valid = np.where(mask > 0.5)[0]
            return int(np.random.choice(valid)) if len(valid) > 0 else 255

        elif spec.agent_type == "greedy":
            return int(game.greedy_action())

        elif spec.agent_type == "mcts":
            sims = spec.config.get("simulations", 100)
            return int(game.mcts_action(sims))

        elif spec.agent_type == "alphabeta":
            depth = spec.config.get("depth", 6)
            return int(game.alphabeta_action(depth))

        else:
            # Neural agent - use the agent's select_action
            return agent.select_action(obs, mask)

    def evaluate(
        self,
        agent_specs: list[str],
        games_per_matchup: int | None = None,
    ) -> EvaluationResults:
        """Evaluate agents against each other.

        Args:
            agent_specs: List of agent specification strings
            games_per_matchup: Override default games per matchup

        Returns:
            EvaluationResults with all metrics
        """
        games_per = games_per_matchup or self.games_per_matchup
        start_time = time.time()

        # Check availability
        available, missing = self.check_availability(agent_specs)

        if missing:
            self._log("Missing agents:")
            for spec in missing:
                self._log(f"  - {spec.name}")
            self._log("")

            # Prompt user if in interactive mode
            if self.verbose:
                self._log("Continue without missing agents? [Y/n] ")
                # For non-interactive, just continue
                pass

        if not available:
            raise ValueError("No agents available for evaluation")

        self._log(f"Evaluating {len(available)} agents:")
        for spec in available:
            self._log(f"  - {spec.name} ({spec.agent_type})")
        self._log("")

        # Load agents
        agents = {spec.name: (self._load_agent(spec), spec) for spec in available}

        # Run all pairwise matchups
        matchups = []
        total_games = 0
        agent_names = list(agents.keys())

        for i, name1 in enumerate(agent_names):
            for name2 in agent_names[i + 1 :]:
                agent1, spec1 = agents[name1]
                agent2, spec2 = agents[name2]

                self._log(f"{name1} vs {name2}...")

                wins1, wins2, draws = 0, 0, 0

                for game_idx in range(games_per):
                    result = self._play_game(agent1, agent2, spec1, spec2, seed=game_idx + 1000)

                    if result == 0:
                        wins1 += 1
                    elif result == 1:
                        wins2 += 1
                    else:
                        draws += 1

                    # Update in-memory ELO
                    self._update_elo(name1, name2, result)

                matchups.append(
                    MatchResult(
                        agent1=name1,
                        agent2=name2,
                        games=games_per,
                        agent1_wins=wins1,
                        agent2_wins=wins2,
                        draws=draws,
                    )
                )

                total_games += games_per
                self._log(f"  {name1}: {wins1} ({wins1/games_per:.1%}) | {name2}: {wins2} ({wins2/games_per:.1%})")

        # Get final ELO ratings
        elo_ratings = {name: self._elo_ratings.get(name, 1500.0) for name in agent_names}

        elapsed = time.time() - start_time

        return EvaluationResults(
            agents=agent_names,
            matchups=matchups,
            elo_ratings=elo_ratings,
            total_games=total_games,
            elapsed_seconds=elapsed,
        )


# Preset configurations
PRESETS = {
    "validate": {
        "agents": ["random", "greedy"],
        "games_per_matchup": 20,
        "description": "Quick validation with basic bots",
    },
    "benchmark": {
        "agents": ["random", "greedy", "mcts-50", "mcts-100", "alphabeta-6"],
        "games_per_matchup": 100,
        "description": "Full benchmark against standard baselines",
    },
    "neural": {
        "agents": ["ppo-generalist", "greedy", "mcts-100", "alphabeta-6"],
        "games_per_matchup": 50,
        "description": "Evaluate neural agent against strong baselines",
    },
}
