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, step
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DISCARD,
    DRAW_DECK,
    HAND_SIZE,
    LOC_P0_HAND,
    N_ACTIONS,
    N_CARDS,
    PLAY,
)
from tests.lost_cities_jax.conftest import game_count
from tests.lost_cities_jax.helpers import shuffled_order

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


@pytest.mark.slow
def test_legal_mask_matches_slow_bruteforce_checker(pytestconfig):
    state_count = game_count(
        pytestconfig,
        env_name="LOST_CITIES_JAX_MASK_STATES",
        local=1_000,
        ci=10_000,
        full=10_000,
    )
    rng = random.Random(31337)
    state = reset_from_order(jnp.asarray(shuffled_order(30_000_000), dtype=jnp.int8))
    next_seed = 30_000_001

    for _ in range(state_count):
        jax_mask = np.asarray(JIT_MASK(state), dtype=bool)
        slow_mask = _slow_mask(state)
        np.testing.assert_array_equal(jax_mask, slow_mask)

        if bool(state.done):
            state = reset_from_order(jnp.asarray(shuffled_order(next_seed), dtype=jnp.int8))
            next_seed += 1
            continue

        legal_actions = np.flatnonzero(jax_mask)
        action = int(legal_actions[rng.randrange(len(legal_actions))])
        state, _, _ = JIT_STEP(state, jnp.int32(action))


def _slow_mask(state) -> np.ndarray:
    mask = np.zeros(N_ACTIONS, dtype=bool)
    if bool(state.done):
        return mask

    loc = np.asarray(state.card_loc)
    pile = np.asarray(state.pile)
    pile_len = np.asarray(state.pile_len)
    col_top = np.asarray(state.col_top)
    player = int(state.to_move)
    just_discarded = int(state.just_discarded)
    hand_loc = LOC_P0_HAND + player
    hand = sorted(card for card in range(N_CARDS) if loc[card] == hand_loc)

    for hand_slot in range(HAND_SIZE):
        if hand_slot >= len(hand):
            continue
        card = hand[hand_slot]
        color = card // CARDS_PER_COLOR
        slot = card % CARDS_PER_COLOR
        is_handshake = slot < 3
        rank = slot - 1

        for place_type in (PLAY, DISCARD):
            if place_type == PLAY:
                place_ok = (
                    col_top[player, color] == 0 if is_handshake else rank > col_top[player, color]
                )
            else:
                place_ok = True
            if not place_ok:
                continue

            for draw_source in range(6):
                if draw_source == DRAW_DECK:
                    draw_ok = True
                else:
                    src = draw_source - 1
                    same_discard_pile = place_type == DISCARD and src == color
                    after_len = int(pile_len[src]) + int(same_discard_pile)
                    if after_len == 0:
                        draw_ok = False
                    else:
                        after_top = card if same_discard_pile else int(pile[src, after_len - 1])
                        just = card if place_type == DISCARD else just_discarded
                        draw_ok = after_top != just
                mask[hand_slot * 12 + place_type * 6 + draw_source] = draw_ok

    return mask
