from __future__ import annotations

import jax
import jax.numpy as jnp
import numpy as np

from lost_cities_jax import legal_action_mask, reset
from lost_cities_jax.opponents import (
    discard_only_action,
    heuristic_balanced_action,
    heuristic_cautious_action,
)
from lost_cities_jax.ppo import (
    JaxPPOConfig,
    NetworkConfig,
    OpponentConfig,
    PPOHyperConfig,
    RunConfig,
    create_train_state,
    evaluate,
    load_config,
    make_train_iteration,
    policy_by_name,
    random_rollout,
    train,
)


def tiny_config(tmp_path) -> JaxPPOConfig:
    return JaxPPOConfig(
        run=RunConfig(
            experiment_name="pytest-jax-ppo",
            seed=123,
            total_updates=1,
            checkpoint_every=1,
            artifact_root=str(tmp_path),
        ),
        opponent=OpponentConfig(name="discard_only"),
        network=NetworkConfig(hidden_size=32, num_layers=1),
        ppo=PPOHyperConfig(batch_games=8, rollout_steps=16, epochs=1, minibatches=2),
    )


def test_static_opponents_return_legal_actions():
    state = reset(jax.random.PRNGKey(0))
    mask = np.asarray(legal_action_mask(state), dtype=bool)
    for fn in [discard_only_action, heuristic_balanced_action, heuristic_cautious_action]:
        action = int(fn(state, jnp.int32(0), jax.random.PRNGKey(1)))
        assert mask[action]


def test_random_rollout_smoke(tmp_path):
    row = random_rollout(tiny_config(tmp_path))
    assert row["env_steps"] == 8 * 16
    assert 0.0 <= row["play_action_rate"] <= 1.0
    assert row["game_length_mean"] > 0.0


def test_train_checkpoint_and_eval_smoke(tmp_path):
    cfg = tiny_config(tmp_path)
    run_dir = train(cfg)
    assert (run_dir / "latest").exists()
    result = evaluate(cfg, run_dir / "latest", games=8, duplicate=True)
    assert result["games"] == 16
    assert 0.0 <= result["win_rate"] <= 1.0


def test_one_jitted_train_iteration_shapes(tmp_path):
    cfg = tiny_config(tmp_path)
    train_state = create_train_state(cfg, jax.random.PRNGKey(0))
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(jax.random.PRNGKey(1), 8))
    train_iteration = make_train_iteration(cfg, policy_by_name("discard_only"))

    train_state, env_state, rng, metrics = train_iteration(
        train_state, env_state, jax.random.PRNGKey(2), jnp.asarray(1.0)
    )

    assert env_state.to_move.shape == (8,)
    assert "loss" in metrics
    assert jnp.isfinite(metrics["loss"])


def test_load_config_file():
    cfg = load_config("configs/jax_ppo/smoke.yaml")
    assert cfg.opponent.name == "discard_only"
    assert cfg.ppo.batch_games == 64