"""Independent pure-Python Lost Cities rules reference.

This module intentionally uses ordinary Python containers rather than the JAX
state representation. It is deterministic under an explicit ``deck_order`` and
the shared flat action encoding.
"""

from __future__ import annotations

from bisect import insort
from dataclasses import dataclass
from random import Random

N_PLAYERS = 2
N_COLORS = 5
CARDS_PER_COLOR = 12
N_CARDS = N_COLORS * CARDS_PER_COLOR
HAND_SIZE = 8
INITIAL_DEAL = N_PLAYERS * HAND_SIZE
MAX_STEPS = 400
N_ACTIONS = HAND_SIZE * 2 * 6

LOC_DECK = 0
LOC_P0_HAND = 1
LOC_P1_HAND = 2
LOC_P0_BOARD = 3
LOC_P1_BOARD = 4
LOC_DISCARD = 5

NO_CARD = -1
PLAY = 0
DISCARD = 1
DRAW_DECK = 0


@dataclass
class RefState:
    deck_order: list[int]
    draw_ptr: int
    card_loc: list[int]
    hand_public: list[bool]
    hands: list[list[int]]
    board: list[list[list[int]]]
    col_top: list[list[int]]
    piles: list[list[int]]
    to_move: int
    just_discarded: int
    step_count: int
    done: bool


def reset(seed: int | None = None) -> RefState:
    rng = Random(seed)
    order = list(range(N_CARDS))
    rng.shuffle(order)
    return reset_from_order(order)


def reset_from_order(deck_order: list[int] | tuple[int, ...]) -> RefState:
    order = [int(card) for card in deck_order]
    card_loc = [LOC_DECK] * N_CARDS
    p0_hand = sorted(order[:HAND_SIZE])
    p1_hand = sorted(order[HAND_SIZE:INITIAL_DEAL])
    for card in p0_hand:
        card_loc[card] = LOC_P0_HAND
    for card in p1_hand:
        card_loc[card] = LOC_P1_HAND
    return RefState(
        deck_order=order,
        draw_ptr=INITIAL_DEAL,
        card_loc=card_loc,
        hand_public=[False] * N_CARDS,
        hands=[p0_hand, p1_hand],
        board=[[[] for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
        col_top=[[0 for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
        piles=[[] for _ in range(N_COLORS)],
        to_move=0,
        just_discarded=NO_CARD,
        step_count=0,
        done=False,
    )


def clone_state(state: RefState) -> RefState:
    return RefState(
        deck_order=list(state.deck_order),
        draw_ptr=state.draw_ptr,
        card_loc=list(state.card_loc),
        hand_public=list(state.hand_public),
        hands=[list(hand) for hand in state.hands],
        board=[[list(col) for col in player] for player in state.board],
        col_top=[list(player) for player in state.col_top],
        piles=[list(pile) for pile in state.piles],
        to_move=state.to_move,
        just_discarded=state.just_discarded,
        step_count=state.step_count,
        done=state.done,
    )


def decode_action(action: int) -> tuple[int, int, int]:
    action = int(action)
    hand_slot = action // 12
    rem = action % 12
    place_type = rem // 6
    draw_source = rem % 6
    return hand_slot, place_type, draw_source


def hand_cards(state: RefState, player: int | None = None) -> list[int]:
    if player is None:
        player = state.to_move
    return list(state.hands[player])


def legal_action_mask(state: RefState) -> list[bool]:
    bits = legal_action_bits(state)
    return [bool(bits & (1 << action)) for action in range(N_ACTIONS)]


def legal_action_bits(state: RefState) -> int:
    bits = 0
    if state.done:
        return bits

    player = state.to_move
    hand = state.hands[player]
    draw_bits = 1
    for color, pile in enumerate(state.piles):
        if pile and pile[-1] != state.just_discarded:
            draw_bits |= 1 << (color + 1)

    for hand_slot, card in enumerate(hand[:HAND_SIZE]):
        base = hand_slot * 12
        if can_play_card(state, player, card):
            bits |= draw_bits << base

        color = card_color(card)
        discard_bits = draw_bits & ~(1 << (color + 1))
        bits |= discard_bits << (base + 6)
    return bits


def nth_legal_action(bits: int, index: int) -> int:
    """Return the ``index``-th set action bit from low to high."""

    remaining = int(index)
    action = 0
    while bits:
        if bits & 1:
            if remaining == 0:
                return action
            remaining -= 1
        action += 1
        bits >>= 1
    raise IndexError(index)


def step(
    state: RefState, action: int, *, validate: bool = True
) -> tuple[RefState, list[float], bool]:
    if state.done or action < 0 or action >= N_ACTIONS:
        return clone_state(state), [0.0, 0.0], state.done
    if validate and not legal_action_mask(state)[action]:
        return clone_state(state), [0.0, 0.0], state.done

    next_state = clone_state(state)
    _, reward, done = step_in_place(next_state, action, validate=False)
    return next_state, reward, done


def step_in_place(
    state: RefState, action: int, *, validate: bool = True
) -> tuple[RefState, list[float], bool]:
    if state.done or action < 0 or action >= N_ACTIONS:
        return state, [0.0, 0.0], state.done
    if validate and not legal_action_mask(state)[action]:
        return state, [0.0, 0.0], state.done

    player = state.to_move
    hand_slot, place_type, draw_source = decode_action(action)
    hand = state.hands[player]
    card = hand.pop(hand_slot)
    color = card_color(card)

    state.hand_public[card] = False
    if place_type == PLAY:
        state.board[player][color].append(card)
        if not is_handshake(card):
            state.col_top[player][color] = rank(card)
        state.card_loc[card] = LOC_P0_BOARD + player
    else:
        state.piles[color].append(card)
        state.card_loc[card] = LOC_DISCARD

    if draw_source == DRAW_DECK:
        drawn = state.deck_order[state.draw_ptr]
        state.draw_ptr += 1
        public = False
    else:
        src = draw_source - 1
        drawn = state.piles[src].pop()
        public = True

    insort(hand, drawn)
    state.card_loc[drawn] = LOC_P0_HAND + player
    state.hand_public[drawn] = public
    state.just_discarded = NO_CARD
    state.step_count += 1
    state.done = (draw_source == DRAW_DECK and state.draw_ptr >= N_CARDS) or (
        state.step_count >= MAX_STEPS
    )
    state.to_move = 1 - player

    reward = board_score(state) if state.done else [0.0, 0.0]
    return state, reward, state.done


def can_play_card(state: RefState, player: int, card: int) -> bool:
    top_rank = state.col_top[player][card_color(card)]
    if is_handshake(card):
        return top_rank == 0
    return rank(card) > top_rank


def board_score(state: RefState) -> list[float]:
    return [float(sum(score_column(column) for column in player)) for player in state.board]


def score(state: RefState) -> list[float]:
    return board_score(state)


def score_column(column: list[int]) -> int:
    if not column:
        return 0
    handshakes = sum(1 for card in column if is_handshake(card))
    rank_sum = sum(rank(card) for card in column if not is_handshake(card))
    value = (rank_sum - 20) * (1 + handshakes)
    if len(column) >= 8:
        value += 20
    return value


def card_color(card: int) -> int:
    return int(card) // CARDS_PER_COLOR


def card_slot(card: int) -> int:
    return int(card) % CARDS_PER_COLOR


def is_handshake(card: int) -> bool:
    return card_slot(card) < 3


def rank(card: int) -> int:
    return card_slot(card) - 1


def _can_draw_after_place(state: RefState, card: int, place_type: int, draw_source: int) -> bool:
    if draw_source == DRAW_DECK:
        return True
    src = draw_source - 1
    same_discard_pile = place_type == DISCARD and card_color(card) == src
    after_len = len(state.piles[src]) + int(same_discard_pile)
    if after_len == 0:
        return False
    after_top = card if same_discard_pile else state.piles[src][-1]
    just_discarded = card if place_type == DISCARD else state.just_discarded
    return after_top != just_discarded


__all__ = [
    "N_ACTIONS",
    "RefState",
    "board_score",
    "clone_state",
    "decode_action",
    "hand_cards",
    "legal_action_bits",
    "legal_action_mask",
    "nth_legal_action",
    "reset",
    "reset_from_order",
    "score",
    "step",
    "step_in_place",
]