"""Render human-readable JAX PPO Lost Cities transcripts."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path
from typing import NamedTuple

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

from lost_cities_jax.engine import (
    board_score,
    current_hand_sorted,
    decode_action,
    reset_from_order,
    step,
)
from lost_cities_jax.opponents import policy_by_name
from lost_cities_jax.ppo import checkpoint_policy, load_config, make_shuffle_bank
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DISCARD,
    DRAW_DECK,
    LOC_P0_BOARD,
    MAX_STEPS,
    N_CARDS,
    N_COLORS,
    PLAY,
    State,
)

COLOR_NAMES = ["red", "green", "white", "blue", "yellow"]
STEP_JIT = jax.jit(step)


class PolicySpec(NamedTuple):
    label: str
    fn: object


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--opponent", default="heuristic_cautious")
    parser.add_argument("--shuffle-bank-seed", type=int, default=20260704)
    parser.add_argument("--pairs", type=int, default=10)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    cfg = load_config(args.config)
    agent = PolicySpec("gate3_checkpoint", jax.jit(checkpoint_policy(cfg, args.checkpoint)))
    opponent = PolicySpec(args.opponent, jax.jit(policy_by_name(args.opponent)))
    orders = make_shuffle_bank(args.shuffle_bank_seed, args.pairs)

    sections = []
    summaries = []
    for idx, order in enumerate(orders):
        sections.append(f"\n\n## Pair {idx:02d} / learner seat 0\n")
        text, summary = render_game(order, [agent, opponent], agent_seat=0, game_id=idx * 2)
        sections.append(text)
        summaries.append(summary)

        sections.append(f"\n\n## Pair {idx:02d} / learner seat 1\n")
        text, summary = render_game(order, [opponent, agent], agent_seat=1, game_id=idx * 2 + 1)
        sections.append(text)
        summaries.append(summary)

    output = render_summary(summaries, args) + "".join(sections)
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    Path(args.output).write_text(output, encoding="utf-8")


def render_game(
    deck_order: list[int], policies: list[PolicySpec], *, agent_seat: int, game_id: int
) -> tuple[str, dict]:
    state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
    lines = [f"Game {game_id}: P0={policies[0].label}, P1={policies[1].label}"]
    handshake_events = []
    cautious_openings = []
    cautious_useful_discards = []

    while not bool(state.done):
        player = int(state.to_move)
        policy = policies[player]
        key = jax.random.PRNGKey(game_id * 1000 + int(state.step_count))
        hand = [int(x) for x in current_hand_sorted(state, player)]
        action = int(policy.fn(state, jnp.asarray(player, dtype=jnp.int32), key))
        hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
        card = hand[hand_slot]
        drawn = peek_drawn_card(state, place_type, draw_source)
        before_col_len = int(state.col_len[player, card_color(card)])
        before_hand = [card_label(x) for x in hand if x < N_CARDS]
        before_scores = [int(x) for x in board_score(state)]

        if policy.label == "heuristic_cautious":
            if place_type == PLAY and before_col_len == 0:
                cautious_openings.append(
                    {
                        "ply": int(state.step_count),
                        "seat": player,
                        "card": card_label(card),
                        "rank": card_rank(card),
                    }
                )
            if place_type == DISCARD and can_play_card(state, 1 - player, card):
                cautious_useful_discards.append(
                    {
                        "ply": int(state.step_count),
                        "seat": player,
                        "card": card_label(card),
                        "opponent": 1 - player,
                    }
                )

        if place_type == PLAY and is_handshake(card):
            handshake_events.append(
                {
                    "ply": int(state.step_count),
                    "seat": player,
                    "policy": policy.label,
                    "card": card_label(card),
                    "hand": before_hand,
                }
            )

        next_state, _, _ = STEP_JIT(state, jnp.asarray(action, dtype=jnp.int32))
        place_text = "PLAY" if place_type == PLAY else "DISCARD"
        target = "expedition" if place_type == PLAY else "discard"
        draw_text = "deck" if draw_source == DRAW_DECK else f"{COLOR_NAMES[draw_source - 1]} pile"
        lines.append(
            f"Ply {int(state.step_count):03d} P{player} {policy.label}: "
            f"{place_text} {card_label(card)} to {target}; "
            f"DRAW {draw_text} -> {card_label(drawn)}; "
            f"score_before={before_scores}"
        )
        lines.extend(render_board(next_state))
        state = next_state

    final_scores = [int(x) for x in board_score(state)]
    lines.append(f"Final score: P0={final_scores[0]} P1={final_scores[1]}")
    lines.extend(render_final_breakdown(state))
    agent_score = final_scores[agent_seat]
    opponent_score = final_scores[1 - agent_seat]
    summary = {
        "game_id": game_id,
        "agent_seat": agent_seat,
        "agent_score_diff": agent_score - opponent_score,
        "agent_opened_colors": int(jnp.sum(state.col_len[agent_seat] > 0)),
        "length": int(state.step_count),
        "max_steps": int(state.step_count) >= MAX_STEPS,
        "handshake_events": handshake_events,
        "cautious_openings": cautious_openings,
        "cautious_useful_discards": cautious_useful_discards,
    }
    return "\n".join(lines) + "\n", summary


def render_summary(summaries: list[dict], args: argparse.Namespace) -> str:
    diffs = np.asarray([item["agent_score_diff"] for item in summaries], dtype=np.float64)
    opened = np.asarray([item["agent_opened_colors"] for item in summaries], dtype=np.float64)
    lengths = np.asarray([item["length"] for item in summaries], dtype=np.float64)
    handshakes = [event for item in summaries for event in item["handshake_events"]]
    cautious_openings = [event for item in summaries for event in item["cautious_openings"]]
    useful_discards = [event for item in summaries for event in item["cautious_useful_discards"]]
    opening_ranks = Counter(str(event["rank"]) for event in cautious_openings)
    low_openings = [
        event for event in cautious_openings if event["rank"] not in ("HS",) and event["rank"] < 7
    ]

    lines = [
        "# Gate-3 Transcript Dump",
        "",
        f"config: `{args.config}`",
        f"checkpoint: `{args.checkpoint}`",
        f"opponent: `{args.opponent}`",
        f"duplicate pairs: `{args.pairs}`",
        "",
        "## Summary",
        "",
        f"games: {len(summaries)}",
        f"agent_score_diff_mean: {float(np.mean(diffs)):.3f}",
        f"agent_score_diff_min_max: {int(np.min(diffs))} / {int(np.max(diffs))}",
        f"agent_opened_colors_mean: {float(np.mean(opened)):.3f}",
        f"game_length_mean: {float(np.mean(lengths)):.3f}",
        f"game_length_min_p50_p95_max: {int(np.min(lengths))} / "
        f"{float(np.quantile(lengths, 0.50)):.1f} / "
        f"{float(np.quantile(lengths, 0.95)):.1f} / {int(np.max(lengths))}",
        f"max_steps_rate: {float(np.mean([item['max_steps'] for item in summaries])):.3f}",
        f"handshake_play_events: {len(handshakes)}",
        f"cautious_openings: {len(cautious_openings)}",
        f"cautious_opening_rank_counts: {json.dumps(dict(sorted(opening_ranks.items())))}",
        f"cautious_low_openings_lt7: {len(low_openings)}",
        f"cautious_discards_immediately_playable_by_opponent: {len(useful_discards)}",
        "",
        "## Handshake Play Contexts",
        "",
    ]
    if handshakes:
        for event in handshakes[:80]:
            lines.append(
                f"- game_event ply={event['ply']} seat=P{event['seat']} "
                f"policy={event['policy']} card={event['card']} hand={event['hand']}"
            )
        if len(handshakes) > 80:
            lines.append(f"- ... {len(handshakes) - 80} more")
    else:
        lines.append("- none")

    lines.extend(["", "## Cautious Opening Audit", ""])
    if cautious_openings:
        for event in cautious_openings[:80]:
            lines.append(
                f"- ply={event['ply']} seat=P{event['seat']} card={event['card']} "
                f"rank={event['rank']}"
            )
        if len(cautious_openings) > 80:
            lines.append(f"- ... {len(cautious_openings) - 80} more")
    else:
        lines.append("- none")

    lines.extend(["", "## Cautious Useful Discard Audit", ""])
    if useful_discards:
        for event in useful_discards[:80]:
            lines.append(
                f"- ply={event['ply']} seat=P{event['seat']} card={event['card']} "
                f"opponent=P{event['opponent']}"
            )
        if len(useful_discards) > 80:
            lines.append(f"- ... {len(useful_discards) - 80} more")
    else:
        lines.append("- none")
    lines.append("")
    return "\n".join(lines)


def render_board(state: State) -> list[str]:
    scores = [int(x) for x in board_score(state)]
    lines = [f"  Board scores: P0={scores[0]} P1={scores[1]}"]
    for player in range(2):
        parts = []
        for color in range(N_COLORS):
            cards = board_cards(state, player, color)
            score = color_score(cards)
            parts.append(f"{COLOR_NAMES[color]}={format_cards(cards)}({score:+d})")
        lines.append(f"  P{player}: " + " | ".join(parts))
    return lines


def render_final_breakdown(state: State) -> list[str]:
    lines = ["Final color breakdown:"]
    for player in range(2):
        parts = []
        for color in range(N_COLORS):
            cards = board_cards(state, player, color)
            parts.append(f"{COLOR_NAMES[color]} {format_cards(cards)} => {color_score(cards):+d}")
        lines.append(f"  P{player}: " + "; ".join(parts))
    return lines


def board_cards(state: State, player: int, color: int) -> list[int]:
    loc = np.asarray(state.card_loc)
    board_loc = LOC_P0_BOARD + player
    start = color * CARDS_PER_COLOR
    cards = [card for card in range(start, start + CARDS_PER_COLOR) if int(loc[card]) == board_loc]
    return sorted(
        cards, key=lambda card: (0 if is_handshake(card) else 1, card_rank_value(card), card)
    )


def color_score(cards: list[int]) -> int:
    if not cards:
        return 0
    handshakes = sum(1 for card in cards if is_handshake(card))
    ranks = [card_rank_value(card) for card in cards if not is_handshake(card)]
    score = (sum(ranks) - 20) * (1 + handshakes)
    if len(cards) >= 8:
        score += 20
    return score


def peek_drawn_card(state: State, place_type: int, draw_source: int) -> int:
    del place_type
    if draw_source == DRAW_DECK:
        return int(state.deck_order[int(state.draw_ptr)])
    color = draw_source - 1
    length = int(state.pile_len[color])
    return int(state.pile[color, length - 1])


def can_play_card(state: State, player: int, card: int) -> bool:
    top = int(state.col_top[player, card_color(card)])
    if is_handshake(card):
        return top == 0
    return card_rank_value(card) > top


def format_cards(cards: list[int]) -> str:
    if not cards:
        return "[]"
    return "[" + ",".join(card_label(card) for card in cards) + "]"


def card_label(card: int) -> str:
    if card < 0 or card >= N_CARDS:
        return "none"
    color = COLOR_NAMES[card_color(card)][0].upper()
    if is_handshake(card):
        return f"{color}HS{card % CARDS_PER_COLOR + 1}"
    return f"{color}{card_rank_value(card)}"


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


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


def card_rank(card: int) -> int | str:
    return "HS" if is_handshake(card) else card_rank_value(card)


def card_rank_value(card: int) -> int:
    return card % CARDS_PER_COLOR - 1


if __name__ == "__main__":
    main()