"""Adapter from the classic two-phase GUI to a JAX PPO checkpoint."""

from __future__ import annotations

import logging
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

import jax.numpy as jnp
import numpy as np

from lost_cities_jax.human_play import (
    evaluate_agent_policy,
    infer_config_path,
    load_agent,
)
from lost_cities_jax.ppo import load_config
from lost_cities_jax.types import (
    LOC_DECK,
    LOC_DISCARD,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    MAX_PILE_SIZE,
    N_CARDS,
    N_COLORS,
    NO_CARD,
    State,
)

from .game import Card, GameState
from .snapshots import Snapshot, snapshot_from_state


def _jax_card_id(card: Card, handshake_copy: int = 0) -> int:
    slot = handshake_copy if card.is_handshake else card.numeric_value(2) + 1
    return card.color * 12 + slot


PublicHandCounts = Sequence[Mapping[tuple[int, int], int]]


def snapshot_to_jax_state(
    snapshot: Snapshot,
    public_hand_counts: PublicHandCounts | None = None,
) -> tuple[State, list[int]]:
    """Convert a classic card-phase snapshot into the equivalent JAX state.

    The returned list maps each JAX sorted hand slot back to the classic hand
    slot. Handshake copies are interchangeable to both rules engines and the
    policy observation, so they are assigned stable per-color identities here.
    """

    if (
        snapshot.config.n_colors != N_COLORS
        or snapshot.config.n_ranks != 9
        or snapshot.config.n_handshakes != 3
        or snapshot.config.hand_size != 8
        or snapshot.config.min_rank != 2
    ):
        raise ValueError("JAX PPO checkpoints require the classic 5-color, 60-card rules")
    if snapshot.phase != "card":
        raise ValueError("JAX PPO policy conversion must start during the card phase")

    handshake_next = [0] * N_COLORS

    def allocate(card: Card) -> int:
        if not card.is_handshake:
            return _jax_card_id(card)
        copy = handshake_next[card.color]
        if copy >= 3:
            raise ValueError(f"too many handshake copies for color {card.color}")
        handshake_next[card.color] += 1
        return _jax_card_id(card, copy)

    card_loc = np.full(N_CARDS, LOC_DECK, dtype=np.int8)
    hand_public = np.zeros(N_CARDS, dtype=bool)
    zone_ids: list[int] = []
    hand_pairs: list[tuple[int, int]] = []

    deck_ids = [allocate(card) for card in snapshot.deck]
    zone_ids.extend(deck_ids)
    for player, hand in enumerate(snapshot.hands):
        public_remaining = dict(public_hand_counts[player]) if public_hand_counts else {}
        for classic_slot, card in enumerate(hand):
            card_id = allocate(card)
            card_loc[card_id] = LOC_P0_HAND + player
            zone_ids.append(card_id)
            face = (card.color, 0 if card.is_handshake else card.numeric_value(2))
            if public_remaining.get(face, 0) > 0:
                hand_public[card_id] = True
                public_remaining[face] -= 1
            if player == snapshot.current_player:
                hand_pairs.append((card_id, classic_slot))

    col_top = np.zeros((2, N_COLORS), dtype=np.int8)
    col_hs = np.zeros((2, N_COLORS), dtype=np.int8)
    col_len = np.zeros((2, N_COLORS), dtype=np.int8)
    for player, expeditions in enumerate(snapshot.expeditions):
        for color, expedition in enumerate(expeditions):
            for card in expedition:
                card_id = allocate(card)
                card_loc[card_id] = LOC_P0_BOARD + player
                zone_ids.append(card_id)
                col_len[player, color] += 1
                if card.is_handshake:
                    col_hs[player, color] += 1
                else:
                    col_top[player, color] = card.numeric_value(2)

    pile = np.full((N_COLORS, MAX_PILE_SIZE), NO_CARD, dtype=np.int8)
    pile_len = np.zeros(N_COLORS, dtype=np.int8)
    for color, discard in enumerate(snapshot.discards):
        for index, card in enumerate(discard):
            card_id = allocate(card)
            card_loc[card_id] = LOC_DISCARD
            pile[color, index] = card_id
            pile_len[color] += 1
            zone_ids.append(card_id)

    if len(zone_ids) != N_CARDS or len(set(zone_ids)) != N_CARDS:
        raise ValueError("classic snapshot does not contain one complete 60-card deck")

    # Classic draws from the end of its deck list. JAX draws forward from
    # draw_ptr, so its remaining suffix is the reversed classic deck.
    remaining = list(reversed(deck_ids))
    deck_id_set = set(deck_ids)
    used = [card_id for card_id in range(N_CARDS) if card_id not in deck_id_set]
    deck_order = np.asarray([*used, *remaining], dtype=np.int8)
    draw_ptr = len(used)
    hand_slot_map = [slot for _, slot in sorted(hand_pairs)]
    state = State(
        deck_order=jnp.asarray(deck_order),
        draw_ptr=jnp.asarray(draw_ptr, dtype=jnp.int32),
        card_loc=jnp.asarray(card_loc),
        hand_public=jnp.asarray(hand_public),
        col_top=jnp.asarray(col_top),
        col_hs=jnp.asarray(col_hs),
        col_len=jnp.asarray(col_len),
        pile=jnp.asarray(pile),
        pile_len=jnp.asarray(pile_len),
        to_move=jnp.asarray(snapshot.current_player, dtype=jnp.int8),
        just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8),
        step_count=jnp.asarray(snapshot.turn_count, dtype=jnp.int32),
        done=jnp.asarray(snapshot.terminal),
    )
    return state, hand_slot_map


class JaxPPOPolicy:
    """Expose an atomic-action JAX policy through the classic policy protocol."""

    def __init__(self, checkpoint: str | Path, *, config: str | Path | None = None):
        # The classic GUI enables root DEBUG logging for its own interaction
        # trace. Keep JAX/Orbax internals from flooding the launcher terminal.
        for logger_name in ("jax", "absl", "orbax"):
            logging.getLogger(logger_name).setLevel(logging.WARNING)
        cfg_path = infer_config_path(checkpoint) if config is None else Path(config)
        self.cfg = load_config(cfg_path)
        self.params, self.model = load_agent(self.cfg, checkpoint)
        self.checkpoint = Path(checkpoint)
        self.pending_draw: int | None = None
        self.last_evaluation: Any | None = None

    def act(
        self,
        obs_or_state: Any,
        public_hand_counts: PublicHandCounts | None = None,
    ) -> int:
        if not isinstance(obs_or_state, GameState):
            raise TypeError("JAX PPO GUI policy requires a GameState")
        if obs_or_state.phase == "draw":
            if self.pending_draw is None:
                raise RuntimeError("JAX PPO draw requested without a pending atomic action")
            draw = self.pending_draw
            self.pending_draw = None
            return draw

        snapshot = snapshot_from_state(obs_or_state)
        jax_state, hand_slot_map = snapshot_to_jax_state(snapshot, public_hand_counts)
        result = evaluate_agent_policy(
            self.cfg,
            self.params,
            self.model,
            jax_state,
            snapshot.current_player,
        )
        self.last_evaluation = result
        jax_hand_slot = result.action // 12
        remainder = result.action % 12
        place = remainder // 6
        self.pending_draw = remainder % 6
        return 2 * hand_slot_map[jax_hand_slot] + place


__all__ = ["JaxPPOPolicy", "PublicHandCounts", "snapshot_to_jax_state"]
