from __future__ import annotations

import numpy as np
from coolrl_lost_cities.games.classic.game import GameState, classic_config

from coolrl_lost_cities.games.classic.jax_ppo_policy import (
    JaxPPOPolicy,
    snapshot_to_jax_state,
)
from coolrl_lost_cities.games.classic.snapshots import snapshot_from_state
from lost_cities_jax.engine import board_score, current_hand_sorted, legal_action_mask
from lost_cities_jax.human_play import PolicyEval


def test_snapshot_to_jax_state_preserves_initial_public_state() -> None:
    state = GameState.new_game(classic_config(), seed=7)
    snapshot = snapshot_from_state(state)

    converted, hand_slot_map = snapshot_to_jax_state(snapshot)

    assert int(converted.to_move) == snapshot.current_player
    assert int(converted.draw_ptr) == 16
    assert hand_slot_map == sorted(
        hand_slot_map,
        key=lambda slot: (
            snapshot.hands[0][slot].color,
            snapshot.hands[0][slot].rank,
        ),
    )
    assert np.asarray(current_hand_sorted(converted)).shape == (8,)
    assert np.asarray(legal_action_mask(converted)).any()
    assert np.asarray(board_score(converted)).tolist() == [0.0, 0.0]


def test_jax_policy_splits_atomic_action_across_classic_phases(monkeypatch) -> None:
    state = GameState.new_game(classic_config(), seed=11)
    converted, hand_slot_map = snapshot_to_jax_state(snapshot_from_state(state))
    legal = np.flatnonzero(np.asarray(legal_action_mask(converted), dtype=bool))
    atomic_action = int(legal[0])
    expected_slot = hand_slot_map[atomic_action // 12]
    expected_place = (atomic_action % 12) // 6
    expected_draw = atomic_action % 6

    monkeypatch.setattr(
        "coolrl_lost_cities.games.classic.jax_ppo_policy.evaluate_agent_policy",
        lambda *_args: PolicyEval(action=atomic_action, top3=[], value=0.0),
    )
    policy = object.__new__(JaxPPOPolicy)
    policy.cfg = object()
    policy.params = object()
    policy.model = object()
    policy.pending_draw = None
    policy.last_evaluation = None

    card_action = policy.act(state)
    assert card_action == 2 * expected_slot + expected_place
    state.apply_action(card_action)
    assert state.phase == "draw"
    assert policy.act(state) == expected_draw
    assert policy.pending_draw is None