"""Pure JAX rules engine for one round of two-player Lost Cities."""

from __future__ import annotations

import jax
import jax.numpy as jnp

from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DISCARD,
    DRAW_DECK,
    HAND_SIZE,
    INITIAL_DEAL,
    LOC_DECK,
    LOC_DISCARD,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    MAX_PILE_SIZE,
    MAX_STEPS,
    N_ACTIONS,
    N_CARDS,
    N_COLORS,
    NO_CARD,
    PLAY,
    UNKNOWN_CARD,
    State,
)


def reset(rng: jax.Array) -> State:
    """Return a freshly shuffled initial state."""

    deck_order = jax.random.permutation(rng, jnp.arange(N_CARDS, dtype=jnp.int32), axis=0).astype(
        jnp.int8
    )
    return reset_from_order(deck_order)


def reset_from_order(deck_order: jax.Array, first_player: jax.Array | int = 0) -> State:
    """Return an initial state using an explicit 60-card deck permutation.

    The first eight cards are dealt to player 0 and the next eight to player 1.
    The remaining cards are drawn in array order starting at ``draw_ptr == 16``.

    ``first_player`` moves first. Rounds two and three of a classic match are led
    by whoever is ahead on points, so the match layer sets this per round.
    """

    deck_order = jnp.asarray(deck_order, dtype=jnp.int8)
    card_loc = jnp.full((N_CARDS,), LOC_DECK, dtype=jnp.int8)
    p0_hand = deck_order[:HAND_SIZE].astype(jnp.int32)
    p1_hand = deck_order[HAND_SIZE:INITIAL_DEAL].astype(jnp.int32)
    card_loc = card_loc.at[p0_hand].set(jnp.int8(LOC_P0_HAND))
    card_loc = card_loc.at[p1_hand].set(jnp.int8(LOC_P0_HAND + 1))

    return State(
        deck_order=deck_order,
        draw_ptr=jnp.asarray(INITIAL_DEAL, dtype=jnp.int32),
        card_loc=card_loc,
        hand_public=jnp.zeros((N_CARDS,), dtype=jnp.bool_),
        col_top=jnp.zeros((2, N_COLORS), dtype=jnp.int8),
        col_hs=jnp.zeros((2, N_COLORS), dtype=jnp.int8),
        col_len=jnp.zeros((2, N_COLORS), dtype=jnp.int8),
        pile=jnp.full((N_COLORS, MAX_PILE_SIZE), NO_CARD, dtype=jnp.int8),
        pile_len=jnp.zeros((N_COLORS,), dtype=jnp.int8),
        to_move=jnp.asarray(first_player, dtype=jnp.int8),
        just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8),
        step_count=jnp.asarray(0, dtype=jnp.int32),
        done=jnp.asarray(False),
    )


def decode_action(action: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]:
    action = jnp.asarray(action, dtype=jnp.int32)
    hand_slot = action // (2 * 6)
    rem = action % (2 * 6)
    place_type = rem // 6
    draw_source = rem % 6
    return hand_slot, place_type, draw_source


def current_hand_sorted(state: State, player: jax.Array | int | None = None) -> jax.Array:
    """Return the selected player's hand as sorted card ids, padded with 60."""

    if player is None:
        player = state.to_move
    player = jnp.asarray(player, dtype=jnp.int32)
    ids = jnp.arange(N_CARDS, dtype=jnp.int32)
    hand_loc = LOC_P0_HAND + player
    sortable = jnp.where(state.card_loc.astype(jnp.int32) == hand_loc, ids, UNKNOWN_CARD)
    return jnp.sort(sortable)[:HAND_SIZE]


def legal_action_mask(state: State) -> jax.Array:
    """Return the 96-bit legal action mask for ``state.to_move``.

    Illegal actions are masked out before placement and draw are considered as a
    single atomic action. In particular, discarding to a color and then drawing
    from that same discard pile is illegal because the top card would be the
    just-discarded card.
    """

    player = state.to_move.astype(jnp.int32)
    cards = current_hand_sorted(state, player).astype(jnp.int32)
    valid_card = cards < N_CARDS
    color = cards // CARDS_PER_COLOR
    safe_color = jnp.clip(color, 0, N_COLORS - 1)
    slot = cards % CARDS_PER_COLOR
    is_handshake = slot < 3
    rank = jnp.where(is_handshake, 0, slot - 1)

    top_for_card = state.col_top[player, safe_color].astype(jnp.int32)
    play_legal = valid_card & jnp.where(is_handshake, top_for_card == 0, rank > top_for_card)
    discard_legal = valid_card
    place_legal = jnp.stack([play_legal, discard_legal], axis=1)

    src_colors = jnp.arange(N_COLORS, dtype=jnp.int32)
    pre_len = state.pile_len.astype(jnp.int32)
    pre_top_idx = jnp.maximum(pre_len - 1, 0)
    pre_top = jnp.where(
        pre_len > 0,
        state.pile[src_colors, pre_top_idx].astype(jnp.int32),
        NO_CARD,
    )

    place_types = jnp.arange(2, dtype=jnp.int32)
    is_discard_place = place_types == DISCARD
    same_discard_pile = is_discard_place[None, :, None] & (
        safe_color[:, None, None] == src_colors[None, None, :]
    )
    after_len = pre_len[None, None, :] + same_discard_pile.astype(jnp.int32)
    after_top = jnp.where(same_discard_pile, cards[:, None, None], pre_top[None, None, :])
    just_discarded = jnp.where(
        is_discard_place[None, :, None],
        cards[:, None, None],
        state.just_discarded.astype(jnp.int32),
    )
    pile_draw_legal = (after_len > 0) & (after_top != just_discarded)
    deck_draw_legal = jnp.ones((HAND_SIZE, 2, 1), dtype=jnp.bool_)
    draw_legal = jnp.concatenate([deck_draw_legal, pile_draw_legal], axis=2)

    mask = (place_legal[:, :, None] & draw_legal).reshape((N_ACTIONS,))
    return mask & ~state.done


def step(state: State, action: jax.Array) -> tuple[State, jax.Array, jax.Array]:
    """Apply one atomic place+draw action.

    Illegal actions are defined as a no-op with zero reward. Done states are
    also no-op states, which keeps vmapped batches stable when games finish at
    different plies. Training code is expected to prevent illegal actions by
    sampling from ``legal_action_mask``.
    """

    action = jnp.asarray(action, dtype=jnp.int32)
    in_range = (action >= 0) & (action < N_ACTIONS)
    safe_action = jnp.clip(action, 0, N_ACTIONS - 1)
    mask = legal_action_mask(state)
    is_legal = in_range & mask[safe_action]

    proposed = _apply_legal_step(state, safe_action)
    next_state = jax.tree_util.tree_map(
        lambda new, old: jnp.where(is_legal, new, old), proposed, state
    )
    terminal_reward = board_score(proposed)
    reward = jnp.where(
        is_legal & proposed.done, terminal_reward, jnp.zeros((2,), dtype=jnp.float32)
    )
    return next_state, reward, next_state.done


def _apply_legal_step(state: State, action: jax.Array) -> State:
    hand_slot, place_type, draw_source = decode_action(action)
    player = state.to_move.astype(jnp.int32)
    next_player = (1 - player).astype(jnp.int8)
    hand_loc = (LOC_P0_HAND + player).astype(jnp.int8)
    board_loc = (LOC_P0_BOARD + player).astype(jnp.int8)

    hand = current_hand_sorted(state, player)
    raw_card = hand[hand_slot].astype(jnp.int32)
    card = jnp.clip(raw_card, 0, N_CARDS - 1)
    color = card // CARDS_PER_COLOR
    slot = card % CARDS_PER_COLOR
    is_handshake = slot < 3
    rank = (slot - 1).astype(jnp.int8)
    is_play = place_type == PLAY
    is_discard = place_type == DISCARD

    placed_loc = jnp.where(is_play, board_loc, jnp.int8(LOC_DISCARD)).astype(jnp.int8)
    card_loc = state.card_loc.at[card].set(placed_loc)
    hand_public = state.hand_public.at[card].set(False)

    play_delta = is_play.astype(jnp.int8)
    col_len = state.col_len.at[player, color].add(play_delta)
    col_hs = state.col_hs.at[player, color].add((is_play & is_handshake).astype(jnp.int8))
    old_top = state.col_top[player, color]
    new_top = jnp.where(is_play & ~is_handshake, rank, old_top).astype(jnp.int8)
    col_top = state.col_top.at[player, color].set(new_top)

    pile_insert_idx = state.pile_len[color].astype(jnp.int32)
    pile_old_value = state.pile[color, pile_insert_idx]
    pile_new_value = jnp.where(is_discard, card.astype(jnp.int8), pile_old_value).astype(jnp.int8)
    pile = state.pile.at[color, pile_insert_idx].set(pile_new_value)
    pile_len = state.pile_len.at[color].add(is_discard.astype(jnp.int8))

    is_deck_draw = draw_source == DRAW_DECK
    draw_src = jnp.clip(draw_source - 1, 0, N_COLORS - 1)
    safe_ptr = jnp.minimum(state.draw_ptr, N_CARDS - 1)
    deck_card = state.deck_order[safe_ptr].astype(jnp.int32)
    pile_draw_len = pile_len[draw_src].astype(jnp.int32)
    pile_top_idx = jnp.maximum(pile_draw_len - 1, 0)
    pile_card = jnp.clip(pile[draw_src, pile_top_idx].astype(jnp.int32), 0, N_CARDS - 1)
    drawn_card = jnp.where(is_deck_draw, deck_card, pile_card).astype(jnp.int32)

    card_loc = card_loc.at[drawn_card].set(hand_loc)
    hand_public = hand_public.at[drawn_card].set(~is_deck_draw)

    drawn_pile_value = pile[draw_src, pile_top_idx]
    pile = pile.at[draw_src, pile_top_idx].set(
        jnp.where(is_deck_draw, drawn_pile_value, jnp.int8(NO_CARD)).astype(jnp.int8)
    )
    pile_len = pile_len.at[draw_src].add(jnp.where(is_deck_draw, 0, -1).astype(jnp.int8))

    draw_ptr = state.draw_ptr + jnp.where(is_deck_draw, 1, 0).astype(jnp.int32)
    step_count = state.step_count + 1
    natural_done = is_deck_draw & (draw_ptr >= N_CARDS)
    forced_done = step_count >= MAX_STEPS
    done = natural_done | forced_done

    return State(
        deck_order=state.deck_order,
        draw_ptr=draw_ptr,
        card_loc=card_loc,
        hand_public=hand_public,
        col_top=col_top,
        col_hs=col_hs,
        col_len=col_len,
        pile=pile,
        pile_len=pile_len,
        to_move=next_player,
        just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8),
        step_count=step_count,
        done=done,
    )


def board_score(state: State) -> jax.Array:
    """Score the currently played boards for both players."""

    card_ids = jnp.arange(N_CARDS, dtype=jnp.int32)
    colors = card_ids // CARDS_PER_COLOR
    slots = card_ids % CARDS_PER_COLOR
    ranks = jnp.where(slots >= 3, slots - 1, 0)

    player_scores = []
    for player in range(2):
        board_loc = LOC_P0_BOARD + player
        on_board = state.card_loc.astype(jnp.int32) == board_loc
        rank_sums = []
        for color in range(N_COLORS):
            rank_sums.append(jnp.sum(jnp.where(on_board & (colors == color), ranks, 0)))
        rank_sum = jnp.stack(rank_sums).astype(jnp.int32)
        length = state.col_len[player].astype(jnp.int32)
        handshakes = state.col_hs[player].astype(jnp.int32)
        base = (rank_sum - 20) * (1 + handshakes)
        bonus = jnp.where(length >= 8, 20, 0)
        color_score = jnp.where(length > 0, base + bonus, 0)
        player_scores.append(jnp.sum(color_score))

    return jnp.stack(player_scores).astype(jnp.float32)


def score(state: State) -> jax.Array:
    """Final round score. This intentionally matches ``board_score``."""

    return board_score(state)


__all__ = [
    "current_hand_sorted",
    "decode_action",
    "legal_action_mask",
    "reset",
    "reset_from_order",
    "step",
    "score",
    "board_score",
]
