"""Terminal human-play client and log summarizer for JAX PPO checkpoints."""

from __future__ import annotations

import json
import math
import time
import uuid
from collections import defaultdict
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import jax
import jax.numpy as jnp
import numpy as np

from lost_cities_jax.engine import (
    board_score,
    current_hand_sorted,
    decode_action,
    legal_action_mask,
    reset_from_order,
    step,
)
from lost_cities_jax.obs import observation
from lost_cities_jax.ppo import (
    ActorCritic,
    JaxPPOConfig,
    create_train_state,
    load_config,
    mask_logits,
    restore_checkpoint,
)
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DISCARD,
    DRAW_DECK,
    LOC_P0_BOARD,
    MAX_STEPS,
    N_ACTIONS,
    N_CARDS,
    N_COLORS,
    PLAY,
    State,
)

COLOR_NAMES = ("R", "G", "W", "B", "Y")
COLOR_WORDS = {
    "R": 0,
    "RED": 0,
    "G": 1,
    "GREEN": 1,
    "W": 2,
    "WHITE": 2,
    "B": 3,
    "BLUE": 3,
    "Y": 4,
    "YELLOW": 4,
}
DEFAULT_HUMAN_LOG_DIR = Path("/mnt/2tbhdd/coolrl-lost-cities-artifacts/human-play")
DEFAULT_HUMAN_BANK_STATE = DEFAULT_HUMAN_LOG_DIR / "shuffle_bank_state.json"
DEFAULT_HUMAN_BANK_SEED = 20260706


@dataclass(frozen=True)
class PolicyEval:
    action: int
    top3: list[dict[str, Any]]
    value: float


def play_cli(
    *,
    checkpoint: str | Path,
    config: str | Path | None = None,
    seat: int = 0,
    seed: int = DEFAULT_HUMAN_BANK_SEED,
    duplicate: bool = False,
    log_dir: str | Path = DEFAULT_HUMAN_LOG_DIR,
    bank_state: str | Path = DEFAULT_HUMAN_BANK_STATE,
    input_fn: Callable[[str], str] = input,
    output_fn: Callable[[str], None] = print,
) -> list[dict[str, Any]]:
    cfg = load_config(infer_config_path(checkpoint) if config is None else config)
    agent = load_agent(cfg, checkpoint)
    log_dir = Path(log_dir)
    deck_index, deck_order = reserve_deck_order(seed, Path(bank_state))
    session_id = time.strftime("%Y%m%d-%H%M%S") + "-" + uuid.uuid4().hex[:8]

    seats = session_human_seats(seat, duplicate)

    records = []
    set_diff = 0.0
    for game_index, human_seat in enumerate(seats, start=1):
        output_fn("")
        output_fn(f"Game {game_index}/{len(seats)} - human seat P{human_seat}")
        record = play_one_game(
            cfg=cfg,
            agent=agent,
            deck_order=deck_order,
            human_seat=human_seat,
            session_id=session_id,
            duplicate_set_id=session_id if duplicate else None,
            duplicate_game_index=game_index if duplicate else None,
            deck_seed=seed,
            deck_index=deck_index,
            input_fn=input_fn,
            output_fn=output_fn,
        )
        records.append(record)
        set_diff += float(record["human_score_diff"])
        append_human_log(log_dir, record)

    if duplicate:
        output_fn("")
        output_fn(f"Duplicate set total human score diff: {set_diff:+.0f}")
    return records


def summarize_logs(log_dir: str | Path = DEFAULT_HUMAN_LOG_DIR) -> dict[str, Any]:
    records = read_human_logs(log_dir)
    games = len(records)
    duplicate_sets: dict[str, list[dict[str, Any]]] = defaultdict(list)
    singletons = []
    for record in records:
        set_id = record.get("duplicate_set_id")
        if set_id:
            duplicate_sets[set_id].append(record)
        else:
            singletons.append(record)

    set_diffs = [
        sum(float(game["human_score_diff"]) for game in group)
        for group in duplicate_sets.values()
        if len(group) == 2
    ]
    game_diffs = [float(record["human_score_diff"]) for record in records]
    ai_opened = [float(record["ai_opened_colors"]) for record in records]
    lengths = [float(record["game_length"]) for record in records]
    swings = collect_value_swings(records)[:10]
    return {
        "games": games,
        "duplicate_sets": len(set_diffs),
        "single_games": len(singletons),
        "human_game_score_diff_mean": mean(game_diffs),
        "human_game_score_diff_ci95": ci95(game_diffs),
        "human_duplicate_set_diff_mean": mean(set_diffs),
        "human_duplicate_set_diff_ci95": ci95(set_diffs),
        "ai_opened_colors_mean": mean(ai_opened),
        "game_length_mean": mean(lengths),
        "value_swings_top10": swings,
    }


def session_human_seats(seat: int, duplicate: bool) -> list[int]:
    if seat not in {0, 1}:
        raise ValueError("seat must be 0 or 1")
    return [seat, 1 - seat] if duplicate else [seat]


def print_summary(summary: dict[str, Any], output_fn: Callable[[str], None] = print) -> None:
    output_fn(json.dumps(summary, indent=2, sort_keys=True, ensure_ascii=False))


def play_one_game(
    *,
    cfg: JaxPPOConfig,
    agent: tuple[Any, Any],
    deck_order: list[int],
    human_seat: int,
    session_id: str,
    duplicate_set_id: str | None,
    duplicate_game_index: int | None,
    deck_seed: int,
    deck_index: int,
    input_fn: Callable[[str], str],
    output_fn: Callable[[str], None],
    human_action_fn: Callable[[State, int], int] | None = None,
) -> dict[str, Any]:
    params, model = agent
    state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
    actions = []
    move_logs = []

    while not bool(state.done) and int(state.step_count) < MAX_STEPS:
        player = int(state.to_move)
        policy_eval = evaluate_agent_policy(cfg, params, model, state, player)
        output_fn(render_public_state(state, human_seat))
        if player == human_seat:
            if human_action_fn is None:
                action = prompt_human_action(state, human_seat, input_fn, output_fn)
            else:
                action = int(human_action_fn(state, human_seat))
            actor = "human"
        else:
            action = policy_eval.action
            actor = "ai"
            output_fn(f"AI: {describe_action(state, action)}")
        before = state
        state, _, _ = step(state, jnp.asarray(action, dtype=jnp.int32))
        actions.append(int(action))
        move_logs.append(
            {
                "ply": int(before.step_count),
                "actor": actor,
                "player": player,
                "action": int(action),
                "action_text": describe_action(before, action),
                "ai_policy_player": player,
                "ai_top3": policy_eval.top3,
                "ai_value": policy_eval.value,
            }
        )

    output_fn(render_public_state(state, human_seat))
    output_fn(render_score_breakdown(state))
    scores = np.asarray(board_score(state), dtype=np.float32)
    ai_seat = 1 - human_seat
    comment = input_fn("Post-game comment (optional, Enter to skip): ").strip()
    return {
        "schema": "lost-cities-jax-human-play-v1",
        "session_id": session_id,
        "duplicate_set_id": duplicate_set_id,
        "duplicate_game_index": duplicate_game_index,
        "deck_seed": deck_seed,
        "deck_index": deck_index,
        "human_seat": human_seat,
        "ai_seat": ai_seat,
        "deck_order": [int(card) for card in deck_order],
        "actions": actions,
        "moves": move_logs,
        "final_scores": {"p0": float(scores[0]), "p1": float(scores[1])},
        "human_score_diff": float(scores[human_seat] - scores[ai_seat]),
        "ai_score_diff": float(scores[ai_seat] - scores[human_seat]),
        "score_breakdown": score_breakdown(state),
        "game_length": int(state.step_count),
        "ai_opened_colors": int(np.asarray(state.col_len)[ai_seat].astype(bool).sum()),
        "human_opened_colors": int(np.asarray(state.col_len)[human_seat].astype(bool).sum()),
        "max_steps": bool(int(state.step_count) >= MAX_STEPS),
        "human_comment": comment,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }


def load_agent(cfg: JaxPPOConfig, checkpoint: str | Path) -> tuple[Any, ActorCritic]:
    state = create_train_state(cfg, jax.random.PRNGKey(0))
    state = restore_checkpoint(Path(checkpoint), state)
    model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
    return state.params, model


def infer_config_path(checkpoint: str | Path) -> Path:
    path = Path(checkpoint).resolve()
    candidates = []
    for parent in [path, *path.parents]:
        candidates.extend([parent / "main_ppo_config.json", parent / "config.json"])
    for candidate in candidates:
        if candidate.exists():
            return candidate
    msg = f"could not infer config for checkpoint; pass --config explicitly: {checkpoint}"
    raise FileNotFoundError(msg)


def evaluate_agent_policy(
    cfg: JaxPPOConfig,
    params: Any,
    model: ActorCritic,
    state: State,
    player: int,
) -> PolicyEval:
    obs = observation(state, jnp.asarray(player, dtype=jnp.int32))[None, :]
    legal = legal_action_mask(state)[None, :]
    logits, value = model.apply(params, obs)
    masked = mask_logits(logits, legal)
    probs = np.asarray(jax.nn.softmax(masked, axis=-1)[0], dtype=np.float64)
    legal_np = np.asarray(legal[0], dtype=bool)
    order = np.argsort(-probs)
    top = []
    for action in order:
        if not legal_np[action]:
            continue
        top.append(
            {
                "action": int(action),
                "prob": float(probs[action]),
                "text": describe_action(state, int(action)),
            }
        )
        if len(top) == 3:
            break
    return PolicyEval(action=int(order[0]), top3=top, value=float(np.asarray(value)[0]))


def prompt_human_action(
    state: State,
    player: int,
    input_fn: Callable[[str], str],
    output_fn: Callable[[str], None],
) -> int:
    while True:
        raw = input_fn("Your move (e.g. play R7 draw deck): ").strip()
        try:
            action = parse_human_action(raw, state, player)
        except ValueError as exc:
            output_fn(f"Invalid input: {exc}")
            continue
        mask = np.asarray(legal_action_mask(state), dtype=bool)
        if action < 0 or action >= N_ACTIONS or not mask[action]:
            output_fn(f"Illegal move: {illegal_reason(state, action)}")
            continue
        return action


def parse_human_action(raw: str, state: State, player: int) -> int:
    parts = raw.strip().split()
    if len(parts) != 4 or parts[2].lower() != "draw":
        msg = "expected '<play|discard> <card> draw <deck|R|G|W|B|Y>'"
        raise ValueError(msg)
    place_word, card_word, _, draw_word = parts
    place = parse_place(place_word)
    card_id = parse_card_token(card_word)
    hand = [
        int(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS
    ]
    matching = [idx for idx, card in enumerate(hand) if same_card_face(card, card_id)]
    if not matching:
        msg = f"card is not in your hand: {card_word}"
        raise ValueError(msg)
    hand_slot = matching[0]
    draw_source = parse_draw_source(draw_word)
    return encode_action(hand_slot, place, draw_source)


def parse_place(value: str) -> int:
    lowered = value.lower()
    if lowered == "play":
        return PLAY
    if lowered == "discard":
        return DISCARD
    msg = "placement must be 'play' or 'discard'"
    raise ValueError(msg)


def parse_draw_source(value: str) -> int:
    upper = value.upper()
    if upper == "DECK":
        return DRAW_DECK
    if upper in COLOR_WORDS:
        return COLOR_WORDS[upper] + 1
    msg = "draw source must be deck or one of R/G/W/B/Y"
    raise ValueError(msg)


def parse_card_token(value: str) -> int:
    token = value.strip().upper()
    if len(token) < 2:
        raise ValueError("card token is too short")
    color = COLOR_WORDS.get(token[0])
    if color is None:
        raise ValueError("card color must be R/G/W/B/Y")
    suffix = token[1:]
    if suffix in {"H", "HS", "W", "WAGER"}:
        slot = 0
    else:
        try:
            rank = int(suffix)
        except ValueError as exc:
            raise ValueError("rank must be 2..10 or HS") from exc
        if rank < 2 or rank > 10:
            raise ValueError("rank must be 2..10")
        slot = rank + 1
    return color * CARDS_PER_COLOR + slot


def same_card_face(left: int, right: int) -> bool:
    return left // CARDS_PER_COLOR == right // CARDS_PER_COLOR and card_label(left) == card_label(
        right
    )


def encode_action(hand_slot: int, place_type: int, draw_source: int) -> int:
    return hand_slot * 12 + place_type * 6 + draw_source


def illegal_reason(state: State, action: int) -> str:
    if action < 0 or action >= N_ACTIONS:
        return "action id outside action space"
    hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(card) for card in np.asarray(current_hand_sorted(state)) if int(card) < N_CARDS]
    if hand_slot >= len(hand):
        return "selected hand slot is empty"
    card = hand[hand_slot]
    color = card // CARDS_PER_COLOR
    rank = card_rank(card)
    player = int(state.to_move)
    if place_type == PLAY:
        top = int(np.asarray(state.col_top)[player, color])
        if rank == 0 and top > 0:
            return "handshake cannot be played after a number card"
        if rank > 0 and rank <= top:
            return f"number card must be above current top rank {top}"
    if draw_source != DRAW_DECK:
        pile_color = draw_source - 1
        pile_len = int(np.asarray(state.pile_len)[pile_color])
        if pile_len <= 0:
            return "discard pile is empty"
        if place_type == DISCARD and pile_color == color:
            return "cannot draw the card you just discarded"
    return "move is not legal under the current mask"


def render_public_state(state: State, player: int) -> str:
    scores = np.asarray(board_score(state), dtype=np.float32)
    opponent = 1 - player
    lines = [
        "",
        f"Ply {int(state.step_count)} | to move P{int(state.to_move)} | deck left {N_CARDS - int(state.draw_ptr)}",
        f"Board score: P{player} {scores[player]:+.0f} / P{opponent} {scores[opponent]:+.0f} "
        f"(you {scores[player] - scores[opponent]:+.0f})",
        f"Your hand: {' '.join(card_name(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS)}",
        "Boards:",
        *render_boards(state),
        "Discards:",
        *render_piles(state),
    ]
    return "\n".join(lines)


def render_boards(state: State) -> list[str]:
    loc = np.asarray(state.card_loc)
    lines = []
    for player in range(2):
        color_parts = []
        for color in range(N_COLORS):
            cards = [
                card
                for card in range(N_CARDS)
                if loc[card] == LOC_P0_BOARD + player and card // CARDS_PER_COLOR == color
            ]
            cards.sort(key=lambda card: (card_rank(card) > 0, card_rank(card), card))
            color_parts.append(f"{COLOR_NAMES[color]}:[{' '.join(card_name(c) for c in cards)}]")
        lines.append(f"  P{player} " + " ".join(color_parts))
    return lines


def render_piles(state: State) -> list[str]:
    pile = np.asarray(state.pile)
    pile_len = np.asarray(state.pile_len)
    lines = []
    for color in range(N_COLORS):
        cards = [int(card) for card in pile[color, : int(pile_len[color])]]
        lines.append(f"  {COLOR_NAMES[color]}: {' '.join(card_name(card) for card in cards)}")
    return lines


def render_score_breakdown(state: State) -> str:
    breakdown = score_breakdown(state)
    lines = ["Final scoring:"]
    for player in range(2):
        parts = []
        for color in COLOR_NAMES:
            item = breakdown[f"p{player}"][color]
            parts.append(f"{color}:{item['score']:+.0f}")
        lines.append(f"  P{player} {' '.join(parts)} total {breakdown[f'p{player}']['total']:+.0f}")
    return "\n".join(lines)


def score_breakdown(state: State) -> dict[str, Any]:
    loc = np.asarray(state.card_loc)
    col_len = np.asarray(state.col_len)
    col_hs = np.asarray(state.col_hs)
    data: dict[str, Any] = {}
    for player in range(2):
        total = 0
        player_data: dict[str, Any] = {}
        for color in range(N_COLORS):
            ranks = [
                card_rank(card)
                for card in range(N_CARDS)
                if loc[card] == LOC_P0_BOARD + player
                and card // CARDS_PER_COLOR == color
                and card_rank(card) > 0
            ]
            length = int(col_len[player, color])
            hs = int(col_hs[player, color])
            rank_sum = sum(ranks)
            score = 0 if length == 0 else (rank_sum - 20) * (1 + hs) + (20 if length >= 8 else 0)
            total += score
            player_data[COLOR_NAMES[color]] = {
                "rank_sum": rank_sum,
                "handshakes": hs,
                "length": length,
                "score": float(score),
            }
        player_data["total"] = float(total)
        data[f"p{player}"] = player_data
    return data


def describe_action(state: State, action: int) -> str:
    hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
    player = int(state.to_move)
    hand = [
        int(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS
    ]
    card = hand[hand_slot] if hand_slot < len(hand) else -1
    place = "play" if place_type == PLAY else "discard"
    draw = "deck" if draw_source == DRAW_DECK else COLOR_NAMES[draw_source - 1]
    return f"{place} {card_name(card)} draw {draw}"


def card_name(card: int) -> str:
    if card < 0 or card >= N_CARDS:
        return "?"
    return COLOR_NAMES[card // CARDS_PER_COLOR] + card_label(card)


def card_label(card: int) -> str:
    rank = card_rank(card)
    return "HS" if rank == 0 else str(rank)


def card_rank(card: int) -> int:
    slot = card % CARDS_PER_COLOR
    return 0 if slot < 3 else slot - 1


def reserve_deck_order(seed: int, bank_state: Path) -> tuple[int, list[int]]:
    bank_state.parent.mkdir(parents=True, exist_ok=True)
    if bank_state.exists():
        state = json.loads(bank_state.read_text(encoding="utf-8"))
        next_index = int(state.get(str(seed), 0))
    else:
        state = {}
        next_index = 0
    state[str(seed)] = next_index + 1
    bank_state.write_text(json.dumps(state, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return next_index, human_deck_order(seed, next_index)


def human_deck_order(seed: int, index: int) -> list[int]:
    rng = np.random.default_rng(seed + index)
    return rng.permutation(N_CARDS).astype("int8").astype(int).tolist()


def append_human_log(log_dir: Path, record: dict[str, Any]) -> Path:
    log_dir.mkdir(parents=True, exist_ok=True)
    path = log_dir / "games.jsonl"
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
    return path


def read_human_logs(log_dir: str | Path) -> list[dict[str, Any]]:
    path = Path(log_dir) / "games.jsonl"
    if not path.exists():
        return []
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def collect_value_swings(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
    swings = []
    for record in records:
        previous = None
        for move in record.get("moves", []):
            value = float(move.get("ai_value", 0.0))
            if previous is not None:
                swings.append(
                    {
                        "session_id": record.get("session_id"),
                        "duplicate_game_index": record.get("duplicate_game_index"),
                        "ply": move.get("ply"),
                        "delta": value - previous,
                        "abs_delta": abs(value - previous),
                        "action_text": move.get("action_text"),
                    }
                )
            previous = value
    swings.sort(key=lambda row: row["abs_delta"], reverse=True)
    return swings


def mean(values: list[float]) -> float | None:
    if not values:
        return None
    return float(np.mean(np.asarray(values, dtype=np.float64)))


def ci95(values: list[float]) -> list[float | None]:
    if not values:
        return [None, None]
    arr = np.asarray(values, dtype=np.float64)
    if arr.size <= 1:
        value = float(arr[0])
        return [value, value]
    se = float(np.std(arr, ddof=1) / math.sqrt(arr.size))
    avg = float(np.mean(arr))
    return [avg - 1.959963984540054 * se, avg + 1.959963984540054 * se]


__all__ = [
    "append_human_log",
    "card_name",
    "describe_action",
    "human_deck_order",
    "infer_config_path",
    "parse_human_action",
    "play_cli",
    "play_one_game",
    "read_human_logs",
    "render_public_state",
    "reserve_deck_order",
    "session_human_seats",
    "summarize_logs",
]
