"""Game-specific tokenizer for Essence Wars transcripts.

Converts structured game states and actions into token sequences for
transformer training. Vocabulary of ~600 tokens covering cards, actions,
stats, keywords, and game structure.

Design principles:
- Every card gets its own token (semantic identity > raw stats)
- Stats use direct integer tokens (hp:0 through hp:30)
- Actions are structured (type + parameters as separate tokens)
- Board state is variable-length (skip empty slots)
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml  # type: ignore[import-untyped]

# ---------------------------------------------------------------------------
# Card database loader
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class CardInfo:
    """Minimal card metadata for tokenization."""

    id: int
    name: str
    card_type: str  # creature, support, spell, commander
    faction: str  # argentum, symbiote, obsidion, neutral
    cost: int
    attack: int | None  # Only for creatures
    health: int | None  # Only for creatures
    keywords: tuple[str, ...]


def _normalize_name(name: str) -> str:
    """Convert card name to token-safe form: 'Venom Fang' -> 'Venom_Fang'."""
    return name.replace(" ", "_").replace("'", "").replace(",", "")


def load_card_database(data_dir: str | Path | None = None) -> dict[int, CardInfo]:
    """Load all cards from YAML data files.

    Returns dict mapping card_id -> CardInfo.
    """
    if data_dir is None:
        # Find data dir relative to this file (project root / data)
        data_dir = Path(__file__).parent.parent.parent.parent / "data"
    data_dir = Path(data_dir)

    cards: dict[int, CardInfo] = {}

    # Load faction cards
    factions = ["argentum", "symbiote", "obsidion", "neutral"]
    card_types = ["creatures", "supports", "spells"]

    for faction in factions:
        for card_type in card_types:
            yaml_path = data_dir / "cards" / "core_set" / faction / f"{card_type}.yaml"
            if not yaml_path.exists():
                continue

            with yaml_path.open() as f:
                data = yaml.safe_load(f)

            for card in data.get("cards", []):
                info = CardInfo(
                    id=card["id"],
                    name=card["name"],
                    card_type=card["card_type"],
                    faction=faction,
                    cost=card.get("cost", 0),
                    attack=card.get("attack"),
                    health=card.get("health"),
                    keywords=tuple(card.get("keywords", [])),
                )
                cards[info.id] = info

    # Load commanders
    for faction in ["argentum", "symbiote", "obsidion"]:
        yaml_path = data_dir / "commanders" / f"{faction}.yaml"
        if not yaml_path.exists():
            continue

        with yaml_path.open() as f:
            data = yaml.safe_load(f)

        for cmd in data.get("commanders", []):
            info = CardInfo(
                id=cmd["id"],
                name=cmd["name"],
                card_type="commander",
                faction=faction,
                cost=0,
                attack=None,
                health=None,
                keywords=(),
            )
            cards[info.id] = info

    return cards


# ---------------------------------------------------------------------------
# Vocabulary builder
# ---------------------------------------------------------------------------

def build_vocabulary(card_db: dict[int, CardInfo]) -> tuple[dict[str, int], dict[int, str]]:
    """Build token vocabulary from card database.

    Returns (token_to_id, id_to_token) dicts.
    """
    tokens: list[str] = []

    def add(tok: str) -> None:
        if tok not in _seen:
            tokens.append(tok)
            _seen.add(tok)

    _seen: set[str] = set()

    # === Special tokens ===
    for tok in ["[PAD]", "[BOS]", "[EOS]", "[SEP]", "[UNK]", "[MASK]"]:
        add(tok)

    # === Game structure ===
    add("[GAME]")
    add("[TURN]")
    add("[STATE]")
    add("[HAND]")
    add("[BOARD]")
    add("[ACTION]")
    add("[RESULT]")
    add("[LEGAL]")
    add("[VALUE]")
    add("[GAMEOVER]")

    # Player tokens
    add("P1")
    add("P2")

    # Game modes
    add("mode:attrition")
    add("mode:essence_war")

    # Turn numbers (game max is 30)
    for i in range(31):
        add(f"turn:{i}")

    # === Deck name tokens ===
    deck_names = [
        "architect_fortify", "artificer_tokens", "sanctum_healer", "vex_piercing",
        "archon_burst", "deathmaster_assassin", "shadow_weaver", "sovereign_lifesteal",
        "alpha_frenzy", "broodmother_pack", "grove_regenerate", "plague_volatile",
    ]
    for d in deck_names:
        add(f"deck:{d}")

    # === Card name tokens ===
    # Each unique card gets its own token
    for card in sorted(card_db.values(), key=lambda c: c.id):
        add(f"card:{_normalize_name(card.name)}")

    # Unknown card fallback
    add("card:[UNK]")

    # === Card type tokens ===
    for ct in ["creature", "support", "spell", "commander"]:
        add(f"type:{ct}")

    # === Stat value tokens ===
    # Life: 0-30 (commanders have 30 life)
    for i in range(-5, 31):
        add(f"hp:{i}")

    # Attack: 0-15
    for i in range(16):
        add(f"att:{i}")

    # Essence: 0-10
    for i in range(11):
        add(f"ess:{i}")

    # Max essence: 0-10
    for i in range(11):
        add(f"maxess:{i}")

    # Cost: 0-10
    for i in range(11):
        add(f"cost:{i}")

    # Deck/hand size: 0-40
    for i in range(41):
        add(f"deck_size:{i}")
    for i in range(11):
        add(f"hand_size:{i}")

    # Action points: 0-5
    for i in range(6):
        add(f"ap:{i}")

    # Damage amounts: 0-15
    for i in range(16):
        add(f"dmg:{i}")

    # === Keyword tokens ===
    keywords = [
        "Guard", "Piercing", "Shield", "Rush", "Regenerate", "Lethal",
        "Frenzy", "Charge", "Volatile", "Ranged", "Ephemeral", "Stealth",
        "Quick", "Lifesteal", "Fortify", "Ward",
    ]
    for kw in keywords:
        add(f"kw:{kw}")

    # === Action tokens ===
    add("act:PlayCard")
    add("act:Attack")
    add("act:UseAbility")
    add("act:CommanderInsight")
    add("act:EndTurn")

    # Slots (board positions 0-4)
    for i in range(5):
        add(f"slot:{i}")

    # Hand positions (0-9)
    for i in range(10):
        add(f"hand:{i}")

    # Targets
    add("target:none")
    add("target:self")
    for i in range(5):
        add(f"target:enemy_{i}")
    add("target:commander")

    # === Board state tokens ===
    add("empty")
    add("occupied")

    # === Game result tokens ===
    add("winner:P1")
    add("winner:P2")
    add("winner:draw")
    add("reason:life_zero")
    add("reason:deck_out")
    add("reason:turn_limit")

    # === Action-value tokens (discretized into 21 bins: -1.0 to +1.0 in 0.1 steps) ===
    for i in range(-10, 11):
        v = i / 10.0
        add(f"av:{v:.1f}")

    # === State value token (same binning) ===
    for i in range(-10, 11):
        v = i / 10.0
        add(f"sv:{v:.1f}")

    # === Reasoning trace tokens (Chain-of-thought from search) ===
    add("[THINK]")  # Start of reasoning section
    for i in range(1, 11):
        add(f"rank:{i}")  # Move ranking (1=best)
    # Relative evaluation bins: normalized so best=+1.0, worst=-1.0
    for i in range(-10, 11):
        v = i / 10.0
        add(f"eval:{v:.1f}")

    # Build bidirectional mapping
    token_to_id = {tok: idx for idx, tok in enumerate(tokens)}
    id_to_token = dict(enumerate(tokens))

    return token_to_id, id_to_token


# ---------------------------------------------------------------------------
# Tokenizer class
# ---------------------------------------------------------------------------

class EssenceWarsTokenizer:
    """Tokenizer for Essence Wars game transcripts.

    Converts structured game data (dicts) into token ID sequences and back.
    """

    def __init__(self, data_dir: str | Path | None = None) -> None:
        self.card_db = load_card_database(data_dir)
        self.token_to_id, self.id_to_token = build_vocabulary(self.card_db)

        # Reverse lookup: card name -> card info
        self._name_to_card: dict[str, CardInfo] = {}
        for card in self.card_db.values():
            self._name_to_card[_normalize_name(card.name)] = card

        # Card ID -> token name
        self._card_id_to_token: dict[int, str] = {}
        for card in self.card_db.values():
            tok = f"card:{_normalize_name(card.name)}"
            self._card_id_to_token[card.id] = tok

    @property
    def vocab_size(self) -> int:
        return len(self.token_to_id)

    @property
    def pad_token_id(self) -> int:
        return self.token_to_id["[PAD]"]

    @property
    def bos_token_id(self) -> int:
        return self.token_to_id["[BOS]"]

    @property
    def eos_token_id(self) -> int:
        return self.token_to_id["[EOS]"]

    def encode_token(self, token: str) -> int:
        """Encode a single token string to its ID."""
        return self.token_to_id.get(token, self.token_to_id["[UNK]"])

    def decode_token(self, token_id: int) -> str:
        """Decode a single token ID to its string."""
        return self.id_to_token.get(token_id, "[UNK]")

    def encode_tokens(self, tokens: list[str]) -> list[int]:
        """Encode a list of token strings to IDs."""
        return [self.encode_token(t) for t in tokens]

    def decode_tokens(self, token_ids: list[int]) -> list[str]:
        """Decode a list of token IDs to strings."""
        return [self.decode_token(t) for t in token_ids]

    def decode_to_string(self, token_ids: list[int]) -> str:
        """Decode token IDs to a human-readable string."""
        tokens = self.decode_tokens(token_ids)
        # Group by structure tokens for readability
        lines: list[str] = []
        current_line: list[str] = []
        for tok in tokens:
            if tok.startswith("[") and tok != "[PAD]" and current_line:
                lines.append(" ".join(current_line))
                current_line = []
            if tok != "[PAD]":
                current_line.append(tok)
        if current_line:
            lines.append(" ".join(current_line))
        return "\n".join(lines)

    # -------------------------------------------------------------------
    # Game state encoding
    # -------------------------------------------------------------------

    def card_token(self, card_id: int) -> str:
        """Get the token string for a card ID."""
        return self._card_id_to_token.get(card_id, "card:[UNK]")

    def encode_game_header(self, game: dict[str, Any]) -> list[str]:
        """Encode game setup info as tokens."""
        tokens = ["[BOS]", "[GAME]"]
        tokens.append(f"deck:{game.get('deck1', 'unknown')}")
        tokens.append(f"deck:{game.get('deck2', 'unknown')}")
        mode = game.get("mode", "essence_war")
        tokens.append(f"mode:{mode}")

        # Commanders
        if "commander1_id" in game:
            tokens.append(self.card_token(game["commander1_id"]))
        if "commander2_id" in game:
            tokens.append(self.card_token(game["commander2_id"]))

        return tokens

    def encode_state(self, state: dict[str, Any]) -> list[str]:
        """Encode a game state snapshot as tokens.

        Expected state dict keys:
            turn, player, p1_life, p2_life, p1_essence, p1_max_essence,
            p2_essence, p2_max_essence, p1_hand (list of card_ids),
            p2_hand_size, p1_board (list of creature dicts), p2_board,
            p1_deck_size, p2_deck_size, p1_ap
        """
        tokens: list[str] = ["[TURN]"]
        tokens.append(f"turn:{state.get('turn', 0)}")
        tokens.append("P1" if state.get("player", 0) == 0 else "P2")

        # Player 1 stats
        tokens.append("[STATE]")
        tokens.append("P1")
        tokens.append(f"hp:{_clamp(state.get('p1_life', 30), -5, 30)}")
        tokens.append(f"ess:{_clamp(state.get('p1_essence', 0), 0, 10)}")
        tokens.append(f"maxess:{_clamp(state.get('p1_max_essence', 0), 0, 10)}")
        tokens.append(f"ap:{_clamp(state.get('p1_ap', 0), 0, 5)}")
        tokens.append(f"deck_size:{_clamp(state.get('p1_deck_size', 0), 0, 40)}")

        # Player 1 hand
        p1_hand = state.get("p1_hand", [])
        if p1_hand:
            tokens.append("[HAND]")
            tokens.append(f"hand_size:{len(p1_hand)}")
            for card_id in p1_hand:
                tokens.append(self.card_token(card_id))

        # Player 2 stats
        tokens.append("[SEP]")
        tokens.append("P2")
        tokens.append(f"hp:{_clamp(state.get('p2_life', 30), -5, 30)}")
        tokens.append(f"ess:{_clamp(state.get('p2_essence', 0), 0, 10)}")
        tokens.append(f"maxess:{_clamp(state.get('p2_max_essence', 0), 0, 10)}")
        tokens.append(f"deck_size:{_clamp(state.get('p2_deck_size', 0), 0, 40)}")

        # Player 2 hand (only size, we don't see opponent's cards)
        p2_hand_size = state.get("p2_hand_size", 0)
        if p2_hand_size > 0:
            tokens.append(f"hand_size:{p2_hand_size}")

        # Board state - Player 1
        p1_board = state.get("p1_board", [])
        tokens.append("[BOARD]")
        tokens.append("P1")
        if not p1_board:
            tokens.append("empty")
        else:
            for creature in p1_board:
                tokens.extend(self._encode_creature(creature))

        # Board state - Player 2
        p2_board = state.get("p2_board", [])
        tokens.append("P2")
        if not p2_board:
            tokens.append("empty")
        else:
            for creature in p2_board:
                tokens.extend(self._encode_creature(creature))

        return tokens

    def _encode_creature(self, creature: dict[str, Any]) -> list[str]:
        """Encode a single creature on board."""
        tokens: list[str] = []
        tokens.append(f"slot:{creature.get('slot', 0)}")
        tokens.append(self.card_token(creature.get("card_id", 0)))
        tokens.append(f"att:{_clamp(creature.get('attack', 0), 0, 15)}")
        tokens.append(f"hp:{_clamp(creature.get('health', 0), -5, 30)}")

        # Keywords
        for kw in creature.get("keywords", []):
            tok = f"kw:{kw}"
            if tok in self.token_to_id:
                tokens.append(tok)

        return tokens

    def encode_action(self, action: dict[str, Any]) -> list[str]:
        """Encode a game action as tokens.

        Expected action dict keys depend on type:
            PlayCard: card_id, hand_index, slot, cost
            Attack: attacker_slot, defender_slot
            UseAbility: slot, ability_index, target
            EndTurn: (no extra keys)
            CommanderInsight: (no extra keys)
        """
        tokens: list[str] = ["[ACTION]"]
        action_type = action.get("type", "EndTurn")

        if action_type == "PlayCard":
            tokens.append("act:PlayCard")
            tokens.append(f"hand:{action.get('hand_index', 0)}")
            tokens.append(self.card_token(action.get("card_id", 0)))
            tokens.append(f"slot:{action.get('slot', 0)}")
            tokens.append(f"cost:{_clamp(action.get('cost', 0), 0, 10)}")
        elif action_type == "Attack":
            tokens.append("act:Attack")
            tokens.append(f"slot:{action.get('attacker_slot', 0)}")
            tokens.append(f"target:enemy_{action.get('defender_slot', 0)}")
        elif action_type == "UseAbility":
            tokens.append("act:UseAbility")
            tokens.append(f"slot:{action.get('slot', 0)}")
            target = action.get("target", "none")
            if isinstance(target, int):
                tokens.append(f"target:enemy_{target}")
            elif target == "self":
                tokens.append("target:self")
            else:
                tokens.append("target:none")
        elif action_type == "CommanderInsight":
            tokens.append("act:CommanderInsight")
        else:
            tokens.append("act:EndTurn")

        return tokens

    def encode_reasoning_trace(
        self,
        ranked_actions: list[dict[str, Any]],
        max_candidates: int = 5,
    ) -> list[str]:
        """Encode a reasoning trace (chain-of-thought from search).

        Each candidate move gets a rank and normalized evaluation score.
        Scores are normalized relative to the position: best=+1.0, worst=-1.0.

        Args:
            ranked_actions: List of dicts with keys:
                type: action type string
                card_id: (optional) card involved
                slot: (optional) board slot
                attacker_slot, defender_slot: (optional) for attacks
                raw_score: float from AB search
            max_candidates: Max candidates to include (default: 5)

        Returns:
            Token list starting with [THINK]
        """
        if not ranked_actions:
            return []

        tokens: list[str] = ["[THINK]"]

        # Normalize scores: best=+1.0, worst=-1.0
        scores = [a["raw_score"] for a in ranked_actions]
        best = max(scores)
        worst = min(scores)
        score_range = best - worst if best != worst else 1.0

        candidates = ranked_actions[:max_candidates]
        for i, action in enumerate(candidates):
            rank = i + 1
            if rank > 10:
                break
            tokens.append(f"rank:{rank}")

            # Compact action description (same as encode_legal_actions)
            action_type = action.get("type", "EndTurn")
            if action_type == "PlayCard":
                tokens.append("act:PlayCard")
                tokens.append(self.card_token(action.get("card_id", 0)))
                tokens.append(f"slot:{action.get('slot', 0)}")
            elif action_type == "Attack":
                tokens.append("act:Attack")
                tokens.append(f"slot:{action.get('attacker_slot', 0)}")
                tokens.append(f"target:enemy_{action.get('defender_slot', 0)}")
            elif action_type == "UseAbility":
                tokens.append("act:UseAbility")
                tokens.append(f"slot:{action.get('slot', 0)}")
            elif action_type == "CommanderInsight":
                tokens.append("act:CommanderInsight")
            else:
                tokens.append("act:EndTurn")

            # Normalized eval score
            if best == worst:
                norm = 0.0  # All moves equally valued
            else:
                norm = (action["raw_score"] - worst) / score_range  # 0 to 1
                norm = norm * 2.0 - 1.0  # -1 to +1
            norm = round(norm, 1)
            norm = max(-1.0, min(1.0, norm))
            tokens.append(f"eval:{norm:.1f}")

        return tokens

    def encode_legal_actions(self, legal_actions: list[dict[str, Any]]) -> list[str]:
        """Encode the list of legal actions (for action-value prediction)."""
        tokens: list[str] = ["[LEGAL]"]
        for la in legal_actions:
            # Encode each legal action compactly (skip [ACTION] prefix)
            action_type = la.get("type", "EndTurn")
            if action_type == "PlayCard":
                tokens.append("act:PlayCard")
                tokens.append(self.card_token(la.get("card_id", 0)))
                tokens.append(f"slot:{la.get('slot', 0)}")
            elif action_type == "Attack":
                tokens.append("act:Attack")
                tokens.append(f"slot:{la.get('attacker_slot', 0)}")
                tokens.append(f"target:enemy_{la.get('defender_slot', 0)}")
            elif action_type == "UseAbility":
                tokens.append("act:UseAbility")
                tokens.append(f"slot:{la.get('slot', 0)}")
            elif action_type == "CommanderInsight":
                tokens.append("act:CommanderInsight")
            else:
                tokens.append("act:EndTurn")

            # Action value (if annotated)
            if "value" in la:
                av = _clamp(round(la["value"], 1), -1.0, 1.0)
                tokens.append(f"av:{av:.1f}")

            tokens.append("[SEP]")

        return tokens

    def encode_game_result(self, result: dict[str, Any]) -> list[str]:
        """Encode game outcome."""
        tokens: list[str] = ["[GAMEOVER]"]

        winner = result.get("winner")
        if winner == 0:
            tokens.append("winner:P1")
        elif winner == 1:
            tokens.append("winner:P2")
        else:
            tokens.append("winner:draw")

        reason = result.get("reason", "life_zero")
        reason_tok = f"reason:{reason}"
        if reason_tok in self.token_to_id:
            tokens.append(reason_tok)

        tokens.append("[EOS]")
        return tokens

    def encode_position(self, position: dict[str, Any]) -> list[str]:
        """Encode a single position (state + action) for per-position training.

        This is the primary encoding for Phase 6.1: each training example is
        one decision point. The model sees game header + current state and
        predicts the action (+ optionally value/action-values).

        Args:
            position: dict with keys:
                game_header: dict (deck1, deck2, mode, commander IDs)
                state: dict (turn, player, life, essence, hand, board, etc.)
                action: dict (the action taken)
                state_value: float (optional, game outcome from this position)
                legal_actions: list[dict] (optional, with per-action values)

        Returns:
            List of token strings
        """
        tokens: list[str] = []

        # Game context
        tokens.extend(self.encode_game_header(position.get("game_header", {})))
        tokens.append("[SEP]")

        # Current state
        tokens.extend(self.encode_state(position.get("state", {})))

        # State value (if available)
        if "state_value" in position:
            sv = _clamp(round(position["state_value"], 1), -1.0, 1.0)
            tokens.append("[VALUE]")
            tokens.append(f"sv:{sv:.1f}")

        # Reasoning trace from search (chain-of-thought)
        if "reasoning_trace" in position:
            tokens.extend(
                self.encode_reasoning_trace(position["reasoning_trace"])
            )

        # Legal actions with values (if available)
        if "legal_actions" in position:
            tokens.extend(self.encode_legal_actions(position["legal_actions"]))

        # The action taken (prediction target)
        tokens.extend(self.encode_action(position.get("action", {"type": "EndTurn"})))

        return tokens

    def encode_full_game(self, game: dict[str, Any]) -> list[str]:
        """Encode a full game transcript as a token sequence.

        For Phase 6.2+: full trajectory modeling where the model sees
        the entire game history.

        Args:
            game: dict with keys:
                deck1, deck2, mode, commander1_id, commander2_id
                turns: list of dicts, each with state + actions
                result: dict with winner, reason

        Returns:
            List of token strings
        """
        tokens: list[str] = []
        tokens.extend(self.encode_game_header(game))
        tokens.append("[SEP]")

        for turn in game.get("turns", []):
            tokens.extend(self.encode_state(turn.get("state", {})))
            for action in turn.get("actions", []):
                tokens.extend(self.encode_action(action))

        tokens.extend(self.encode_game_result(game.get("result", {})))
        return tokens


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _clamp(value: int | float, lo: int | float, hi: int | float) -> int | float:
    """Clamp value to [lo, hi]."""
    if value < lo:
        return lo
    if value > hi:
        return hi
    return value
