"""Evaluate Generals.io agents by playing them against each other.

Run with `.venv/bin/python scripts/evaluate.py` or `uv run python scripts/evaluate.py`.
Override defaults with OmegaConf arguments such as `num_games=200 agent0.kind=hunter`.

Each agent is either one of the bundled JAX agents (`random`, `expander`,
`hunter`) or a trained policy loaded from an equinox checkpoint via
`agent0.kind=model agent0.checkpoint=ppo_model.eqx`. Games are run in a single
vmapped batch of `num_games` parallel envs.
"""

from dataclasses import dataclass, field

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom
from generals.core import game
from generals.core.action import compute_valid_move_mask
from generals.core.env import GeneralsEnv
from generals.core.observation import Observation
from omegaconf import DictConfig, OmegaConf

from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
from general_bots_training.opponents import OPPONENT_TYPES, StaticOpponent


@dataclass
class AgentConfig:
    """One side of the matchup.

    kind: "random" | "expander" | "hunter" | "model"
    checkpoint: path to an equinox checkpoint (only used when kind == "model")
    """

    kind: str = "random"
    checkpoint: str | None = None


@dataclass
class EvalConfig:
    grid_dims: tuple[int, int] = (21, 21)
    truncation: int = 500
    num_games: int = 100
    agent0: AgentConfig = field(default_factory=lambda: AgentConfig(kind="random"))
    agent1: AgentConfig = field(default_factory=lambda: AgentConfig(kind="expander"))
    seed: int = 0


def load_config(args: list[str] | None = None) -> DictConfig:
    """Merge CLI overrides into the typed default evaluation configuration."""
    defaults = OmegaConf.structured(EvalConfig)
    return OmegaConf.merge(defaults, OmegaConf.from_cli(args))  # type: ignore


class NetworkAgent:
    """Wrap a PolicyValueNetwork in the stateless StaticOpponent interface."""

    def __init__(self, network):
        self._network = network

    def act(self, observation: Observation, key):
        obs_arr = obs_to_tensor(observation)
        mask = compute_valid_move_mask(
            observation.armies, observation.owned_cells, observation.mountains
        )
        action, _, _, _ = self._network(obs_arr, mask, key, None)
        return action


def make_agent(cfg: AgentConfig, key) -> StaticOpponent:
    """Build an agent from its config. `key` seeds the network when needed."""
    kind = cfg.kind.lower().replace("-", "_")
    if kind == "model":
        if cfg.checkpoint is None:
            raise ValueError("agent kind 'model' requires a checkpoint path")
        network = PolicyValueNetwork(key, in_channels=14)
        network = eqx.tree_deserialise_leaves(cfg.checkpoint, network)
        return NetworkAgent(network)
    if kind in OPPONENT_TYPES:
        return OPPONENT_TYPES[kind]()
    choices = ", ".join([*sorted(OPPONENT_TYPES), "model"])
    raise ValueError(f"unknown agent kind {cfg.kind!r}; choose one of: {choices}")


def agent_label(cfg: AgentConfig) -> str:
    return cfg.checkpoint if cfg.kind == "model" and cfg.checkpoint else cfg.kind


def main(config: DictConfig):
    key = jrandom.PRNGKey(config.seed)
    key, agent0_key, agent1_key, env_key = jrandom.split(key, 4)

    agent0 = make_agent(config.agent0, agent0_key)
    agent1 = make_agent(config.agent1, agent1_key)
    label0, label1 = agent_label(config.agent0), agent_label(config.agent1)

    env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
    pool, _ = env.reset(env_key)
    get_obs = game.get_full_observation if env.perfect_info else game.get_observation

    n = config.num_games
    key, init_key = jrandom.split(key)
    init_keys = jrandom.split(init_key, n)
    states = jax.vmap(env.init_state)(init_keys)
    states = states._replace(pool_idx=jnp.arange(n, dtype=states.pool_idx.dtype) % env.pool_size)

    step_env = jax.vmap(env.step, in_axes=(0, 0, None))

    def run_episode(states, key):
        """Play one full game per env; return final winner per env.

        The env auto-resets on done and overwrites `state.winner` with -1, so we
        capture the winner from `timestep.info` on the first done step per env and
        latch it for the remainder of the scan.
        """

        def body(carry, _):
            states, latched_winner, key = carry
            key, p0_key, p1_key = jrandom.split(key, 3)
            obs_p0 = jax.vmap(lambda s: get_obs(s, 0))(states)
            obs_p1 = jax.vmap(lambda s: get_obs(s, 1))(states)
            keys_p0 = jrandom.split(p0_key, n)
            keys_p1 = jrandom.split(p1_key, n)
            actions_p0 = jax.vmap(agent0.act)(obs_p0, keys_p0)
            actions_p1 = jax.vmap(agent1.act)(obs_p1, keys_p1)
            actions = jnp.stack([actions_p0, actions_p1], axis=1)
            timesteps, states = step_env(states, actions, pool)
            done = timesteps.terminated | timesteps.truncated
            # On the first done, latch the winner (truncated games stay at -1 = draw).
            new_winner = jnp.where(
                done & (latched_winner < 0), timesteps.info.winner, latched_winner
            )
            return (states, new_winner, key), done

        # Run a fixed number of steps equal to the truncation length, which
        # guarantees every env has terminated or truncated at least once.
        (states, latched_winner, key), _ = jax.lax.scan(
            body, (states, jnp.full(n, -1, dtype=jnp.int32), key), None, length=config.truncation
        )
        return latched_winner

    print("Generals.io evaluation")
    print(OmegaConf.to_yaml(config, resolve=True).rstrip())
    print(f"device: {jax.devices()[0]}")
    print(f"agent0: {label0}")
    print(f"agent1: {label1}")
    print(f"games:  {n}")
    print()
    print("warming up (jit compile)...")
    winners = run_episode(states, key)
    jax.block_until_ready(winners)
    print("warming up done\n")

    print("playing...")
    winners = run_episode(states, key)
    jax.block_until_ready(winners)

    wins0 = int(jnp.sum(winners == 0))
    wins1 = int(jnp.sum(winners == 1))
    draws = int(jnp.sum(winners < 0))

    print("\nsummary")
    print("-" * 40)
    print(f"agent0 ({label0}): {wins0:4d} wins ({wins0 / n * 100:.1f}%)")
    print(f"agent1 ({label1}): {wins1:4d} wins ({wins1 / n * 100:.1f}%)")
    print(f"draws:              {draws:4d}      ({draws / n * 100:.1f}%)")
    print("-" * 40)
    print(f"total games: {n}")


if __name__ == "__main__":
    main(load_config())