import json
import random

import pytest

from wizard_game import GameOptions, WizardGame

from wizard_ai.agents import HeuristicAgent, ModelAgent, RandomAgent, make_agent
from wizard_ai.model import FEATURE_SIZE, ActionSpace, UniformPolicy, encode_observation, load_policy
from wizard_ai.training import SelfPlayTrainer, TrainingConfig, WizardEnv, evaluate


@pytest.mark.parametrize("num_players", [3, 4, 5, 6])
def test_encoding_has_fixed_size_and_masks_match(num_players):
    game = WizardGame(num_players, GameOptions(hidden_bids=True), seed=num_players)
    rng = random.Random(0)
    while not game.is_finished:
        obs = game.observation(game.current_player)
        assert len(encode_observation(obs)) == FEATURE_SIZE
        mask = ActionSpace.mask(obs)
        legal = {ActionSpace.index(a) for a in obs.legal_actions}
        assert {i for i, allowed in enumerate(mask) if allowed} == legal
        game.apply(obs.player, rng.choice(obs.legal_actions))


def test_action_space_round_trip():
    for index in range(ActionSpace.SIZE):
        assert ActionSpace.index(ActionSpace.action(index)) == index


@pytest.mark.parametrize("agent", [HeuristicAgent(), RandomAgent(1), ModelAgent(UniformPolicy(1))])
def test_agents_play_legal_games(agent):
    game = WizardGame(4, GameOptions(not_equal=True), seed=3)
    while not game.is_finished:
        obs = game.observation(game.current_player)
        action = agent.act(obs)
        assert action in obs.legal_actions
        game.apply(obs.player, action)


def test_env_rewards_add_up_to_scores():
    env = WizardEnv(3, seed=5)
    obs = env.reset()
    totals = [0.0] * env.num_players
    rng = random.Random(5)
    while obs is not None:
        result = env.step(rng.choice(obs.legal_actions))
        totals = [t + r for t, r in zip(totals, result.rewards)]
        obs = result.observation
    assert [round(t * 100) for t in totals] == env.game.scores


def test_trainer_runs_and_checkpoints(tmp_path):
    config = TrainingConfig(episodes=4, player_counts=(3, 4), episodes_per_update=2,
                            checkpoint_dir=tmp_path / "ckpt", seed=1)
    trainer = SelfPlayTrainer(UniformPolicy(1), config)
    trajectories, scores = trainer.play_episode()
    assert len(trajectories) == len(scores)
    assert all(t[-1].done for t in trajectories)
    assert sum(sum(tr.reward for tr in t) for t in trajectories) == pytest.approx(sum(scores) / 100)
    trainer.train()
    assert json.loads((tmp_path / "ckpt" / "policy.json").read_text())["type"] == "uniform"
    assert isinstance(load_policy(tmp_path / "ckpt"), UniformPolicy)


def test_trainer_against_heuristic_opponents():
    config = TrainingConfig(episodes=2, player_counts=(5,), opponents="heuristic", seed=2)
    trajectories, _ = SelfPlayTrainer(UniformPolicy(2), config).play_episode()
    assert len(trajectories) == 1


def test_make_agent_falls_back_without_checkpoint(tmp_path):
    assert isinstance(make_agent("model", tmp_path / "missing"), HeuristicAgent)
    UniformPolicy().save(tmp_path / "ckpt")
    assert isinstance(make_agent("model", tmp_path / "ckpt"), ModelAgent)


def test_heuristic_beats_random():
    result = evaluate([HeuristicAgent(), RandomAgent(1), RandomAgent(2), RandomAgent(3)], games=40, seed=0)
    assert result.mean_scores[0] > max(result.mean_scores[1:])