app.py

"""Tiny Hanabi: Play Hanabi with a trained AI model."""

import json
import random
import re
from dataclasses import dataclass

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Game Configuration
COLORS = ("R", "G")
RANKS = (1, 2, 3)
HAND_SIZE = 2
MAX_INFO_TOKENS = 8
MAX_LIFE_TOKENS = 3
MODEL_ID = "nph4rd/Qwen3-1.7B-Tiny-Hanabi-XML-RL-12-2"

COLOR_NAMES = {"R": "Red", "G": "Green"}
COLOR_HEX = {"R": "#e63946", "G": "#2a9d8f"}
COLOR_HEX_DARK = {"R": "#9d2933", "G": "#1d6e64"}

# Hanabi-themed CSS
CUSTOM_CSS = """
.gradio-container {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f0f23 100%) !important;
    min-height: 100vh;
}

.main-title {
    text-align: center;
    color: #ffd700 !important;
    text-shadow: 0 0 20px #ffd700, 0 0 40px #ff6b6b;
    font-size: 2.5em !important;
    margin-bottom: 0.2em !important;
}

.subtitle {
    text-align: center;
    color: #a0a0a0 !important;
    font-size: 1.1em !important;
}

.game-panel {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 1px solid rgba(255, 215, 0, 0.3) !important;
    border-radius: 12px !important;
    padding: 20px !important;
}

.firework-display {
    background: linear-gradient(180deg, #1a1a2e 0%, #0d0d1a 100%) !important;
    border: 2px solid #ffd700 !important;
    border-radius: 12px !important;
    padding: 15px !important;
    text-align: center;
}

.card-red {
    background: linear-gradient(135deg, #e63946 0%, #9d2933 100%) !important;
    color: white !important;
    padding: 8px 16px !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    display: inline-block !important;
    margin: 4px !important;
    min-width: 50px !important;
    text-align: center !important;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3) !important;
}

.card-green {
    background: linear-gradient(135deg, #2a9d8f 0%, #1d6e64 100%) !important;
    color: white !important;
    padding: 8px 16px !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    display: inline-block !important;
    margin: 4px !important;
    min-width: 50px !important;
    text-align: center !important;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3) !important;
}

.card-unknown {
    background: linear-gradient(135deg, #4a4a6a 0%, #2d2d4a 100%) !important;
    color: #ffd700 !important;
    padding: 8px 16px !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    display: inline-block !important;
    margin: 4px !important;
    min-width: 50px !important;
    text-align: center !important;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3) !important;
    border: 2px dashed #ffd700 !important;
}

.card-empty {
    background: rgba(50, 50, 70, 0.5) !important;
    color: #666 !important;
    padding: 8px 16px !important;
    border-radius: 8px !important;
    display: inline-block !important;
    margin: 4px !important;
    min-width: 50px !important;
    text-align: center !important;
    border: 1px dashed #444 !important;
}

.token-info {
    color: #4dabf7 !important;
    font-size: 1.3em !important;
}

.token-life {
    color: #e63946 !important;
    font-size: 1.3em !important;
}

.score-display {
    color: #ffd700 !important;
    font-size: 1.5em !important;
    font-weight: bold !important;
}

.history-panel {
    background: rgba(0, 0, 0, 0.4) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 8px !important;
    padding: 15px !important;
    max-height: 200px !important;
    overflow-y: auto !important;
    color: #ffffff !important;
}

.history-panel p {
    color: #ffffff !important;
}

.action-btn {
    background: linear-gradient(135deg, #ffd700 0%, #ffaa00 100%) !important;
    color: #1a1a2e !important;
    font-weight: bold !important;
    border: none !important;
    padding: 10px 20px !important;
    border-radius: 8px !important;
}

.action-btn:hover {
    background: linear-gradient(135deg, #ffaa00 0%, #ff8800 100%) !important;
}

.new-game-btn {
    background: linear-gradient(135deg, #2a9d8f 0%, #1d6e64 100%) !important;
    color: white !important;
    font-weight: bold !important;
}

.status-text {
    color: #ffffff !important;
    font-size: 1.1em !important;
    line-height: 1.6 !important;
    padding: 10px !important;
    background: rgba(0, 0, 0, 0.4) !important;
    border-radius: 8px !important;
    margin-top: 10px !important;
}

.status-text p {
    color: #ffffff !important;
}

.game-over-text {
    color: #ffd700 !important;
    font-size: 1.2em !important;
    font-weight: bold !important;
    text-align: center !important;
}
"""


@dataclass
class GameConfig:
    colors: tuple = COLORS
    ranks: tuple = RANKS
    hand_size: int = HAND_SIZE
    max_info_tokens: int = MAX_INFO_TOKENS
    max_life_tokens: int = MAX_LIFE_TOKENS

    @property
    def num_colors(self) -> int:
        return len(self.colors)

    @property
    def max_rank(self) -> int:
        return max(self.ranks)

    @property
    def max_score(self) -> int:
        return self.num_colors * self.max_rank

    def card_distribution(self) -> list[int]:
        """Standard Hanabi distribution: 3 copies of lowest, 2 of middle, 1 of highest."""
        dist = []
        for i, rank in enumerate(self.ranks):
            if i == 0:
                dist.extend([rank, rank, rank])
            elif i == len(self.ranks) - 1:
                dist.append(rank)
            else:
                dist.extend([rank, rank])
        return dist


CONFIG = GameConfig()


def generate_system_prompt(player_id: int) -> str:
    """Generate the system prompt for a player."""
    config = CONFIG
    colors_formatted = ", ".join(f"{c} ({COLOR_NAMES[c]})" for c in config.colors)
    max_rank = config.max_rank
    min_rank = min(config.ranks)
    max_score = config.max_score
    hand_size = config.hand_size
    last_pos = hand_size - 1
    positions = ", ".join(str(i) for i in range(hand_size))

    dist = config.card_distribution()
    cards_per_color = len(dist)
    total_cards = config.num_colors * cards_per_color

    return f"""You are playing Hanabi, a cooperative card game where players work together to build fireworks.

## OBJECTIVE
Build {config.num_colors} fireworks (one per color) by playing cards in sequence from {min_rank} to {max_rank}.
- Colors: {colors_formatted}
- Score: sum of fireworks' highest cards
- Perfect score: {max_score} points (all fireworks at {max_rank})
- Game ends when: all fireworks reach {max_rank}, you lose all {config.max_life_tokens} lives, or final round completes

## FINAL ROUND
When the deck runs out, each player (including the one who drew the last card) gets one final turn. The game ends after the player before the one who drew the last card takes their turn.

## CRITICAL RULE
You CANNOT see your own cards! You can only see other players' hands. You must use hints and deduce what you're holding.

## CARD DISTRIBUTION
The deck has {total_cards} cards total ({cards_per_color} per color):
- Rank 1: 3 copies per color ({3 * config.num_colors} total) - safest to discard duplicates
- Rank 2: 2 copies per color ({2 * config.num_colors} total)
- Rank 3: 1 copy per color ({1 * config.num_colors} total) - NEVER discard, irreplaceable!

## CRITICAL CARDS
- Rank 3s: Only ONE copy of each - if discarded, that color can never reach 3!
- Last copies: Before discarding, check if both copies of a rank 2 are already gone
- Next playable: Cards that are exactly +1 of current firework level are high priority

## STRATEGY PRIORITIES
1. SAFE PLAYS: Play cards you KNOW are playable (confirmed color AND rank match next firework)
2. CRITICAL HINTS: Give hints that enable immediate plays or save critical cards (3s, last copies)
3. SAFE DISCARDS: Discard cards confirmed as unneeded (already played ranks, or known duplicates)
4. RISKY PLAYS: Only when necessary and probability is favorable

## CARD FORMAT
Cards are represented as ColorRank (single letter + number). Examples: R1 (for Red 1), G3 (for Green 3), etc. To represent empty slots (no card), we use --

## YOUR HAND
Your hand has positions {positions} (left to right). You don't know what cards you have unless you've received hints.

What you know about your hand from previous hints is shown as an array, using the following format:

- ?? = unknown color, unknown rank
- C? = known color, unknown rank
- ?R = unknown color, known rank
- CR = known color and known rank

You have full knowledge over other players' hands (e.g., ["R1", "G3"]).

When you play or discard a card, the remaining cards shift left to fill the gap, and the new card is drawn into the rightmost position (position {last_pos}). Hint information shifts with the cards.

## HINT DEDUCTION
When you receive a hint, use it to deduce BOTH positive and negative information:
- Positive: Cards identified by the hint have that color/rank
- Negative: Cards NOT identified by the hint do NOT have that color/rank

## MESSAGE FORMAT
Each turn you'll receive a message with two parts:

1. **Previously** (what happened since your last turn):
   - Player 0: <action result>
   - Player 1: <action result>
   - ...

2. **Current game state** (JSON):
   - info_tokens: Available hint tokens (max {config.max_info_tokens}, gain 1 by discarding)
   - life_tokens: Remaining lives (game over at 0)
   - deck_count: Number of cards remaining in deck
   - fireworks: Object mapping each color to its current level
   - score: Current score (0-{max_score})
   - discards: Array of discarded cards
   - hands: Object with each player's hand (your hand shows hints, others show actual cards)

## AVAILABLE ACTIONS

1. P + position: Play a card from your hand
   - Example: <action>P0</action> (plays card at position 0)
   - If valid (next in sequence), it's added to the firework
   - If invalid, you lose a life and discard the card

2. D + position: Discard a card to gain a hint token
   - Example: <action>D0</action> (discards card at position 0)
   - Gain 1 info token (up to max of {config.max_info_tokens})

3. player + H + color/rank: Give a hint to another player
   - Example: <action>0HR</action> (tells Player 0 which cards are Red)
   - Example: <action>0H2</action> (tells Player 0 which cards are 2s)
   - Costs 1 info token
   - Cannot hint yourself

You are player {player_id}.

Respond with EXACTLY ONE action wrapped in <action></action> tags.
"""


def card_to_str(card: tuple[int, int] | None) -> str:
    """Convert card tuple to string like 'R1' or '--' for empty."""
    if card is None:
        return "--"
    color_idx, rank_idx = card
    return f"{CONFIG.colors[color_idx]}{CONFIG.ranks[rank_idx]}"


class TinyHanabiGame:
    """Manages game state for a Tiny Hanabi game."""

    def __init__(self, seed: int | None = None):
        self.config = CONFIG
        self.seed = seed if seed is not None else random.randint(0, 10000)
        self.reset()

    def reset(self):
        """Reset the game to initial state."""
        random.seed(self.seed)

        # Build and shuffle deck
        dist = self.config.card_distribution()
        self.deck: list[tuple[int, int] | None] = [
            (c, self.config.ranks.index(r))
            for c in range(self.config.num_colors)
            for r in dist
        ]
        random.shuffle(self.deck)

        # Deal hands (2 players)
        self.hands: list[list[tuple[int, int] | None]] = []
        for _ in range(2):
            hand = [self.deck.pop() for _ in range(self.config.hand_size)]
            self.hands.append(hand)

        # Game state
        self.fireworks = {color: 0 for color in self.config.colors}
        self.info_tokens = self.config.max_info_tokens
        self.life_tokens = self.config.max_life_tokens
        self.discard_pile: list[tuple[int, int]] = []
        self.colors_revealed = [[None] * self.config.hand_size for _ in range(2)]
        self.ranks_revealed = [[None] * self.config.hand_size for _ in range(2)]
        self.current_player = 0
        self.score = 0
        self.final_round_turns: int | None = None
        self.game_over = False
        self.game_over_reason: str | None = None
        self.turn_history: list[dict] = []
        # AI conversation history for multi-turn context
        self.ai_messages: list[dict] = [
            {"role": "system", "content": generate_system_prompt(player_id=1)}
        ]

    def get_observation(self, player_id: int) -> str:
        """Generate observation JSON for a player."""
        # Own hand (hints only)
        hand_hints = []
        for idx in range(len(self.hands[player_id])):
            color_hint = self.colors_revealed[player_id][idx]
            rank_hint = self.ranks_revealed[player_id][idx]
            color_str = color_hint if color_hint else "?"
            rank_str = str(rank_hint) if rank_hint else "?"
            hand_hints.append(f"{color_str}{rank_str}")

        # Build hands dict
        hands = {f"player_{player_id}": hand_hints}
        for pid in range(2):
            if pid != player_id:
                cards = [card_to_str(card) for card in self.hands[pid]]
                hands[f"player_{pid}"] = cards

        game_state = {
            "info_tokens": self.info_tokens,
            "life_tokens": self.life_tokens,
            "deck_count": len(self.deck),
            "fireworks": self.fireworks,
            "score": self.score,
            "discards": [card_to_str(card) for card in self.discard_pile],
            "hands": hands,
        }

        if self.game_over:
            game_state["game_over"] = True
            if self.game_over_reason:
                game_state["game_over_reason"] = self.game_over_reason

        # Build output with history
        parts = []
        if self.turn_history:
            parts.append("Previously:")
            for fb in self.turn_history:
                pid = fb["player"]
                msg = fb["feedback"]
                msg_lower = msg[0].lower() + msg[1:] if msg else msg
                parts.append(f"- Player {pid} {msg_lower}")
            parts.append("")
            parts.append("Current game state:")

        parts.append(json.dumps(game_state, indent=2))
        return "\n".join(parts)

    def get_valid_actions(self, player_idx: int) -> list[tuple[str, str]]:
        """Get list of valid actions as (action_code, description) tuples."""
        actions = []
        hand = self.hands[player_idx]

        # Play actions - always valid if card exists
        for pos, card in enumerate(hand):
            if card is not None:
                actions.append((f"P{pos}", f"Play card at position {pos}"))

        # Discard actions - only if below max info tokens
        if self.info_tokens < self.config.max_info_tokens:
            for pos, card in enumerate(hand):
                if card is not None:
                    actions.append((f"D{pos}", f"Discard card at position {pos}"))

        # Hint actions - only if have info tokens
        if self.info_tokens > 0:
            target = 1 - player_idx  # The other player
            target_hand = self.hands[target]

            # Find which colors exist in target's hand
            colors_in_hand = set()
            ranks_in_hand = set()
            for card in target_hand:
                if card is not None:
                    color_idx, rank_idx = card
                    colors_in_hand.add(self.config.colors[color_idx])
                    ranks_in_hand.add(self.config.ranks[rank_idx])

            # Add valid color hints
            for color in colors_in_hand:
                color_name = COLOR_NAMES[color]
                actions.append((f"{target}H{color}", f"Tell AI about {color_name} cards"))

            # Add valid rank hints
            for rank in sorted(ranks_in_hand):
                actions.append((f"{target}H{rank}", f"Tell AI about {rank}s"))

        return actions

    def execute_action(self, player_idx: int, action: str | None) -> str:
        """Execute an action and return feedback."""
        if not action:
            return self._penalize("No action provided.")

        action = action.strip().upper()

        # Play action: P0, P1
        if action.startswith("P") and len(action) >= 2:
            try:
                position = int(action[1])
                return self._play_card(player_idx, position)
            except (ValueError, IndexError):
                return self._penalize(f"Invalid play action '{action}'.")

        # Discard action: D0, D1
        elif action.startswith("D") and len(action) >= 2:
            try:
                position = int(action[1])
                return self._discard_card(player_idx, position)
            except (ValueError, IndexError):
                return self._penalize(f"Invalid discard action '{action}'.")

        # Hint action: 0HR, 1H2, etc.
        elif len(action) >= 3 and action[1] == "H":
            try:
                target = int(action[0])
                hint_value = action[2]
                return self._give_hint(player_idx, target, hint_value)
            except (ValueError, IndexError):
                return self._penalize(f"Invalid hint action '{action}'.")

        else:
            return self._penalize(f"Unknown action '{action}'.")

    def _penalize(self, feedback: str) -> str:
        """Deduct a life for invalid action."""
        self.life_tokens -= 1
        if self.life_tokens <= 0:
            self.game_over = True
            self.game_over_reason = "Lost all lives"
            return f"{feedback} Lost 1 life. Game Over!"
        return f"{feedback} Lost 1 life."

    def _play_card(self, player_idx: int, position: int) -> str:
        """Execute play action."""
        hand = self.hands[player_idx]
        hand_size = len(hand)

        if position < 0 or position >= hand_size:
            return self._penalize(f"Invalid position {position}.")

        card = hand[position]
        if card is None:
            return self._penalize(f"No card at position {position}.")

        color_idx, rank_idx = card
        color = self.config.colors[color_idx]
        rank = self.config.ranks[rank_idx]
        current_level = self.fireworks[color]

        if current_level + 1 == rank:
            # Success
            self.fireworks[color] = rank
            self.score += 1
            feedback = f"Successfully played {card_to_str(card)}!"

            if rank == self.config.max_rank and self.info_tokens < self.config.max_info_tokens:
                self.info_tokens += 1
                feedback += f" +1 info token for completing {COLOR_NAMES[color]}!"

            if self.score == self.config.max_score:
                self.game_over = True
                self.game_over_reason = "Perfect score!"
                feedback += " PERFECT GAME!"
        else:
            # Failure
            self.life_tokens -= 1
            self.discard_pile.append(card)
            expected = current_level + 1
            feedback = f"Played {card_to_str(card)}, but {COLOR_NAMES[color]} needs {expected}. Lost 1 life."

            if self.life_tokens <= 0:
                self.game_over = True
                self.game_over_reason = "Lost all lives"
                return f"{feedback} Game Over!"

        # Shift hand and draw
        self._shift_and_draw(player_idx, position)
        return feedback

    def _discard_card(self, player_idx: int, position: int) -> str:
        """Execute discard action."""
        hand = self.hands[player_idx]
        hand_size = len(hand)

        if position < 0 or position >= hand_size:
            return self._penalize(f"Invalid position {position}.")

        if self.info_tokens >= self.config.max_info_tokens:
            return self._penalize(f"Already at max info tokens ({self.config.max_info_tokens}).")

        card = hand[position]
        if card is None:
            return self._penalize(f"No card at position {position}.")

        self.discard_pile.append(card)
        self.info_tokens += 1
        feedback = f"Discarded {card_to_str(card)}. +1 info token."

        self._shift_and_draw(player_idx, position)
        return feedback

    def _give_hint(self, player_idx: int, target: int, hint_value: str) -> str:
        """Execute hint action."""
        if self.info_tokens <= 0:
            return self._penalize("No info tokens available.")

        if target < 0 or target >= 2:
            return self._penalize(f"Invalid target player {target}.")

        if target == player_idx:
            return self._penalize("Cannot hint yourself.")

        target_hand = self.hands[target]
        matching = []

        if hint_value in self.config.colors:
            color_idx = self.config.colors.index(hint_value)
            for idx, card in enumerate(target_hand):
                if card is not None and card[0] == color_idx:
                    matching.append(idx)
                    self.colors_revealed[target][idx] = hint_value
            hint_type = COLOR_NAMES[hint_value]
        else:
            try:
                hint_num = int(hint_value)
                if hint_num not in self.config.ranks:
                    return self._penalize(f"Invalid rank {hint_num}.")
                rank_idx = self.config.ranks.index(hint_num)
                for idx, card in enumerate(target_hand):
                    if card is not None and card[1] == rank_idx:
                        matching.append(idx)
                        self.ranks_revealed[target][idx] = hint_num
                hint_type = f"{hint_num}s"
            except ValueError:
                return self._penalize(f"Invalid hint value '{hint_value}'.")

        if not matching:
            return self._penalize(f"Player {target} has no {hint_type} cards.")

        self.info_tokens -= 1
        positions_str = ", ".join(str(p) for p in matching)
        target_name = "AI" if target == 1 else "you"
        return f"Hinted {target_name}: {hint_type} at position(s) [{positions_str}]"

    def _shift_and_draw(self, player_idx: int, position: int):
        """Shift cards left and draw new card."""
        hand = self.hands[player_idx]
        hand_size = len(hand)

        for i in range(position, hand_size - 1):
            hand[i] = hand[i + 1]
            self.colors_revealed[player_idx][i] = self.colors_revealed[player_idx][i + 1]
            self.ranks_revealed[player_idx][i] = self.ranks_revealed[player_idx][i + 1]

        if self.deck:
            hand[hand_size - 1] = self.deck.pop()
            if self.final_round_turns is None and len(self.deck) == 0:
                self.final_round_turns = 2  # 2 players
        else:
            hand[hand_size - 1] = None

        self.colors_revealed[player_idx][hand_size - 1] = None
        self.ranks_revealed[player_idx][hand_size - 1] = None

    def advance_turn(self, feedback: str):
        """Record feedback and advance to next player."""
        self.turn_history.append({
            "player": self.current_player,
            "feedback": feedback,
        })

        # Check final round
        if self.final_round_turns is not None:
            self.final_round_turns -= 1
            if self.final_round_turns <= 0:
                self.game_over = True
                self.game_over_reason = "Final round complete"

        self.current_player = 1 - self.current_player


# Global model and tokenizer (loaded once)
model = None
tokenizer = None


def load_model():
    """Load the model and tokenizer."""
    global model, tokenizer
    if model is None:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float32,
            device_map="auto",
        )
        model.eval()
    return model, tokenizer


def get_ai_action(game: TinyHanabiGame) -> str:
    """Get action from the AI model with full conversation history."""
    model, tokenizer = load_model()

    # Get current observation and add to AI's conversation history
    observation = game.get_observation(player_id=1)
    game.ai_messages.append({"role": "user", "content": observation})

    # Generate using full conversation history
    text = tokenizer.apply_chat_template(game.ai_messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

    # Store AI's response in conversation history
    game.ai_messages.append({"role": "assistant", "content": response})

    # Parse action from <action>...</action>
    match = re.search(r"<action>(.*?)</action>", response, re.IGNORECASE | re.DOTALL)
    if match:
        return match.group(1).strip()
    return ""


def format_card_html(card_str: str, is_unknown: bool = False) -> str:
    """Format a card as HTML with appropriate styling."""
    if card_str == "--":
        return '<span class="card-empty">--</span>'

    if is_unknown or "?" in card_str:
        return f'<span class="card-unknown">{card_str}</span>'

    color = card_str[0] if card_str else ""
    if color == "R":
        return f'<span class="card-red">{card_str}</span>'
    elif color == "G":
        return f'<span class="card-green">{card_str}</span>'
    else:
        return f'<span class="card-unknown">{card_str}</span>'


def format_fireworks_html(game: TinyHanabiGame) -> str:
    """Format fireworks display as HTML."""
    parts = []
    for color in CONFIG.colors:
        level = game.fireworks[color]
        color_name = COLOR_NAMES[color]
        color_class = "card-red" if color == "R" else "card-green"

        # Show completed cards
        cards_html = ""
        for rank in range(1, level + 1):
            cards_html += f'<span class="{color_class}" style="margin: 2px;">{color}{rank}</span>'

        # Show remaining slots
        for rank in range(level + 1, CONFIG.max_rank + 1):
            cards_html += f'<span class="card-empty" style="margin: 2px;">{rank}</span>'

        parts.append(f'<div style="margin: 10px 0;"><strong style="color: {COLOR_HEX[color]};">{color_name}:</strong> {cards_html}</div>')

    return "".join(parts)


def format_game_display(game: TinyHanabiGame) -> str:
    """Format the main game display as HTML."""
    if game is None:
        return "<p style='text-align: center; color: #888;'>Start a new game to begin!</p>"

    html = []

    # Score and tokens
    html.append(f'''
    <div style="text-align: center; margin-bottom: 20px;">
        <span class="score-display">Score: {game.score} / {CONFIG.max_score}</span>
    </div>
    ''')

    # Tokens
    info_dots = "🔵 " * game.info_tokens + "⚫ " * (CONFIG.max_info_tokens - game.info_tokens)
    life_dots = "❤️ " * game.life_tokens + "🖤 " * (CONFIG.max_life_tokens - game.life_tokens)

    html.append(f'''
    <div style="display: flex; justify-content: space-around; margin-bottom: 20px;">
        <div><span class="token-info">Info: {info_dots}</span></div>
        <div><span class="token-life">Lives: {life_dots}</span></div>
        <div style="color: #888;">Deck: {len(game.deck)} cards</div>
    </div>
    ''')

    # Fireworks
    html.append('<div class="firework-display">')
    html.append('<h3 style="color: #ffd700; margin-top: 0;">Fireworks</h3>')
    html.append(format_fireworks_html(game))
    html.append('</div>')

    # Hands
    html.append('<div style="margin-top: 20px;">')

    # Your hand (with hints)
    html.append('<div style="margin-bottom: 15px;">')
    html.append('<strong style="color: #ffd700;">Your Hand (Player 0):</strong> ')
    for idx in range(CONFIG.hand_size):
        card = game.hands[0][idx]
        color_hint = game.colors_revealed[0][idx]
        rank_hint = game.ranks_revealed[0][idx]

        if card is None:
            card_str = "--"
        else:
            color_str = color_hint if color_hint else "?"
            rank_str = str(rank_hint) if rank_hint else "?"
            card_str = f"{color_str}{rank_str}"

        html.append(f'<span style="margin-right: 5px;">[{idx}]</span>')
        html.append(format_card_html(card_str, is_unknown=True))

    html.append('</div>')

    # AI hand (fully visible)
    html.append('<div style="margin-bottom: 15px;">')
    html.append('<strong style="color: #4dabf7;">AI Hand (Player 1):</strong> ')
    for idx, card in enumerate(game.hands[1]):
        card_str = card_to_str(card)
        html.append(f'<span style="margin-right: 5px;">[{idx}]</span>')
        html.append(format_card_html(card_str))
    html.append('</div>')

    html.append('</div>')

    # Discards
    if game.discard_pile:
        html.append('<div style="margin-top: 15px; padding: 10px; background: rgba(0,0,0,0.2); border-radius: 8px;">')
        html.append('<strong style="color: #888;">Discards:</strong> ')
        for card in game.discard_pile:
            html.append(format_card_html(card_to_str(card)))
        html.append('</div>')

    return "".join(html)


def format_history(game: TinyHanabiGame) -> str:
    """Format turn history."""
    if game is None or not game.turn_history:
        return "<p style='color: #aaa; font-style: italic;'>No actions yet</p>"

    lines = []
    for entry in game.turn_history:
        player = "You" if entry["player"] == 0 else "AI"
        color = "#ffd700" if entry["player"] == 0 else "#4dabf7"
        lines.append(f'<p style="margin: 5px 0; color: #fff;"><strong style="color: {color};">{player}:</strong> {entry["feedback"]}</p>')

    return "".join(lines)


def new_game(seed_input: str):
    """Start a new game."""
    try:
        seed = int(seed_input) if seed_input.strip() else None
    except ValueError:
        seed = None

    game = TinyHanabiGame(seed=seed)
    valid_actions = game.get_valid_actions(0)
    action_choices = [f"{code}: {desc}" for code, desc in valid_actions]

    status = f'<p style="color: #2a9d8f;">Game started! Seed: {game.seed}</p>'
    status += '<p style="color: #ffd700; font-weight: bold;">Your turn!</p>'

    return (
        game,
        format_game_display(game),
        format_history(game),
        status,
        gr.update(choices=action_choices, value=action_choices[0] if action_choices else None, interactive=True),
        gr.update(interactive=True),
    )


def play_action(game: TinyHanabiGame, action_choice: str):
    """Process human action and AI response."""
    if game is None:
        return (
            None,
            format_game_display(None),
            "",
            '<p style="color: #e63946;">Please start a new game first.</p>',
            gr.update(choices=[], interactive=False),
            gr.update(interactive=False),
        )

    if game.game_over:
        return (
            game,
            format_game_display(game),
            format_history(game),
            f'<p style="color: #ffd700; font-weight: bold;">Game over! {game.game_over_reason}. Final score: {game.score}/{CONFIG.max_score}</p>',
            gr.update(choices=[], interactive=False),
            gr.update(interactive=False),
        )

    # Extract action code from choice string
    action = action_choice.split(":")[0].strip() if action_choice else ""

    # Execute human action
    feedback = game.execute_action(0, action)
    game.advance_turn(feedback)

    if game.game_over:
        status = f'<p><strong style="color: #ffd700;">You:</strong> {feedback}</p>'
        status += f'<p style="color: #ffd700; font-weight: bold; margin-top: 10px;">Game over! {game.game_over_reason}. Final score: {game.score}/{CONFIG.max_score}</p>'
        return (
            game,
            format_game_display(game),
            format_history(game),
            status,
            gr.update(choices=[], interactive=False),
            gr.update(interactive=False),
        )

    # AI turn
    ai_action = get_ai_action(game)
    ai_feedback = game.execute_action(1, ai_action)
    game.advance_turn(ai_feedback)

    status = f'<p><strong style="color: #ffd700;">You:</strong> {feedback}</p>'
    status += f'<p><strong style="color: #4dabf7;">AI ({ai_action}):</strong> {ai_feedback}</p>'

    if game.game_over:
        status += f'<p style="color: #ffd700; font-weight: bold; margin-top: 10px;">Game over! {game.game_over_reason}. Final score: {game.score}/{CONFIG.max_score}</p>'
        return (
            game,
            format_game_display(game),
            format_history(game),
            status,
            gr.update(choices=[], interactive=False),
            gr.update(interactive=False),
        )

    # Update valid actions for next turn
    valid_actions = game.get_valid_actions(0)
    action_choices = [f"{code}: {desc}" for code, desc in valid_actions]

    status += '<p style="color: #2a9d8f; margin-top: 10px;">Your turn!</p>'

    return (
        game,
        format_game_display(game),
        format_history(game),
        status,
        gr.update(choices=action_choices, value=action_choices[0] if action_choices else None, interactive=True),
        gr.update(interactive=True),
    )


# Build UI
with gr.Blocks(title="Tiny Hanabi", css=CUSTOM_CSS, theme=gr.themes.Base()) as demo:
    gr.HTML('<h1 class="main-title">Tiny Hanabi</h1>')
    gr.HTML('<p class="subtitle">Like Hanabi, but tiny: 2 colors, 3 ranks, 2-card hands.</p>')
    gr.HTML('<p class="subtitle" style="font-size: 0.95em; color: #888;">Card distribution per color: 1s (×3), 2s (×2), 3s (×1) — 12 cards total</p>')

    game_state = gr.State(None)

    with gr.Row():
        with gr.Column(scale=2):
            game_display = gr.HTML(
                "<p style='text-align: center; color: #888;'>Start a new game to begin!</p>",
                elem_classes=["game-panel"]
            )

            gr.HTML('<h3 style="color: #ffd700; margin-top: 20px;">History</h3>')
            history_display = gr.HTML("", elem_classes=["history-panel"])

        with gr.Column(scale=1):
            gr.HTML('<h3 style="color: #ffd700;">New Game</h3>')
            with gr.Row():
                seed_input = gr.Textbox(
                    label="Seed (optional)",
                    placeholder="Random",
                    scale=2
                )
                new_game_btn = gr.Button(
                    "New Game",
                    variant="primary",
                    elem_classes=["new-game-btn"],
                    scale=1
                )

            gr.HTML('<h3 style="color: #ffd700; margin-top: 20px;">Your Action</h3>')
            action_dropdown = gr.Dropdown(
                label="Choose Action",
                choices=[],
                interactive=False,
                allow_custom_value=False
            )
            play_btn = gr.Button(
                "Submit Action",
                variant="primary",
                interactive=False,
                elem_classes=["action-btn"]
            )

            status_display = gr.HTML("", elem_classes=["status-text"])

    # Events
    new_game_btn.click(
        new_game,
        inputs=[seed_input],
        outputs=[game_state, game_display, history_display, status_display, action_dropdown, play_btn],
    )

    play_btn.click(
        play_action,
        inputs=[game_state, action_dropdown],
        outputs=[game_state, game_display, history_display, status_display, action_dropdown, play_btn],
    )


if __name__ == "__main__":
    demo.launch()