from __future__ import annotations

import random

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

from lost_cities_jax import legal_action_mask, reset_from_order, score, step
from tests.lost_cities_jax.conftest import game_count
from tests.lost_cities_jax.helpers import (
    assert_state_consistent,
    assert_states_equal,
    shuffled_order,
)

JIT_MASK = jax.jit(legal_action_mask)
JIT_STEP = jax.jit(step)


@pytest.mark.slow
def test_random_trajectories_preserve_state_invariants(pytestconfig):
    games = game_count(
        pytestconfig,
        env_name="LOST_CITIES_JAX_PROPERTY_GAMES",
        local=50,
        ci=1_000,
        full=20_000,
    )
    rng = random.Random(777)

    for game_idx in range(games):
        state = reset_from_order(jnp.asarray(shuffled_order(20_000_000 + game_idx), dtype=jnp.int8))
        assert_state_consistent(state)

        for _ in range(500):
            if bool(state.done):
                after_done, reward, done = JIT_STEP(state, jnp.int32(0))
                assert_states_equal(after_done, state)
                np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
                assert bool(done) is True
                break

            mask = np.asarray(JIT_MASK(state), dtype=bool)
            legal_actions = np.flatnonzero(mask)
            assert len(legal_actions) > 0
            action = int(legal_actions[rng.randrange(len(legal_actions))])
            state, reward, done = JIT_STEP(state, jnp.int32(action))
            assert_state_consistent(state)
            if bool(done):
                np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(state)))
        else:
            pytest.fail(f"game did not terminate within guard loop: game={game_idx}")


def test_forced_termination_scores_current_board():
    state = reset_from_order(jnp.arange(60, dtype=jnp.int8))._replace(
        step_count=jnp.asarray(399, dtype=jnp.int32)
    )
    action = int(np.flatnonzero(np.asarray(legal_action_mask(state), dtype=bool))[0])
    next_state, reward, done = step(state, jnp.int32(action))

    assert bool(done) is True
    assert int(next_state.step_count) == 400
    np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(next_state)))
