"""Generate game transcripts for Game Transformer training.

Uses the Rust game engine via Python bindings to play games between
bots and record structured transcripts that can be tokenized for
transformer training.

Usage:
    uv run python -m essence_wars.transformer.data_generator \
        --games 1000 --bot greedy --output data/transcripts/greedy_1k.jsonl.gz
    uv run python -m essence_wars.transformer.data_generator \
        --games 100 --bot alphabeta3 --with-av --output data/transcripts/ab3_100_av.jsonl.gz
"""

from __future__ import annotations

import gzip
import json
import multiprocessing as mp
import random
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

from essence_wars import PyGame

if TYPE_CHECKING:
    import numpy as np

# ============================================================================
# Constants matching Rust engine (crates/cardgame/src/core/config.rs)
# ============================================================================

STARTING_LIFE = 30
MAX_ESSENCE = 10
AP_PER_TURN = 3
MAX_DECK_SIZE = 60
MAX_HAND_SIZE = 10
CREATURE_SLOTS = 5
SUPPORT_SLOTS = 2
CARD_ID_NORMALIZER = 6000.0
TURN_LIMIT = 30

# Keyword bit definitions (crates/cardgame/src/core/keywords.rs)
KEYWORD_BITS: dict[int, str] = {
    0x0001: "Rush",
    0x0002: "Ranged",
    0x0004: "Piercing",
    0x0008: "Guard",
    0x0010: "Lifesteal",
    0x0020: "Lethal",
    0x0040: "Shield",
    0x0080: "Quick",
    0x0100: "Ephemeral",
    0x0200: "Regenerate",
    0x0400: "Stealth",
    0x0800: "Charge",
    0x1000: "Frenzy",
    0x2000: "Volatile",
    0x4000: "Fortify",
    0x8000: "Ward",
}

# All available decks (3 factions x 4 archetypes)
ALL_DECKS = [
    "architect_fortify", "artificer_tokens", "sanctum_healer", "vex_piercing",
    "archon_burst", "deathmaster_assassin", "shadow_weaver", "sovereign_lifesteal",
    "alpha_frenzy", "broodmother_pack", "grove_regenerate", "plague_volatile",
]


# ============================================================================
# Tensor Decoder
# ============================================================================

def _decode_card_id(normalized: float) -> int:
    """Decode a normalized card ID back to integer."""
    if normalized <= 0:
        return 0
    return round(normalized * CARD_ID_NORMALIZER)


def _decode_keywords(bitfield_normalized: float) -> list[str]:
    """Decode keyword bitfield (normalized 0-1) to keyword names."""
    raw = round(bitfield_normalized * 65535)
    return [name for bit, name in KEYWORD_BITS.items() if raw & bit]


def decode_observation(obs: np.ndarray[Any, Any]) -> dict[str, Any]:
    """Decode 328-float observation tensor into structured game state.

    Returns a dict compatible with EssenceWarsTokenizer.encode_state().
    """
    # Global state (indices 0-5)
    turn = max(1, round(obs[0] * TURN_LIMIT))
    current_player = round(obs[1])

    # Commander IDs (indices 326-327)
    commander1_id = _decode_card_id(obs[326])
    commander2_id = _decode_card_id(obs[327])

    # --- Player 1 (indices 6-80) ---
    p1 = _decode_player_state(obs, offset=6)
    # --- Player 2 (indices 81-155) ---
    p2 = _decode_player_state(obs, offset=81)

    # Decode board creature card IDs from embeddings section (156-325)
    _assign_creature_card_ids(obs, p1, p2)

    return {
        "turn": turn,
        "player": current_player,
        "p1_life": p1["life"],
        "p2_life": p2["life"],
        "p1_essence": p1["essence"],
        "p1_max_essence": min(turn, MAX_ESSENCE),
        "p2_essence": p2["essence"],
        "p2_max_essence": min(turn, MAX_ESSENCE),
        "p1_ap": p1["ap"],
        "p1_deck_size": p1["deck_size"],
        "p2_deck_size": p2["deck_size"],
        "p1_hand": p1["hand_card_ids"],
        "p2_hand_size": p2["hand_size"],
        "p1_board": p1["board"],
        "p2_board": p2["board"],
        # Extra info not used by tokenizer but useful for analysis
        "_p2_hand": p2["hand_card_ids"],
        "_p1_hand_size": p1["hand_size"],
        "_commander1_id": commander1_id,
        "_commander2_id": commander2_id,
    }


def _decode_player_state(obs: np.ndarray[Any, Any], offset: int) -> dict[str, Any]:
    """Decode one player's state from the tensor."""
    life = round(obs[offset + 0] * STARTING_LIFE)
    essence = round(obs[offset + 1] * MAX_ESSENCE)
    ap = round(obs[offset + 2] * AP_PER_TURN)
    deck_size = round(obs[offset + 3] * MAX_DECK_SIZE)
    hand_size = round(obs[offset + 4] * MAX_HAND_SIZE)

    # Hand cards (indices [5-14], normalized card IDs)
    hand_card_ids: list[int] = []
    for i in range(MAX_HAND_SIZE):
        cid = _decode_card_id(obs[offset + 5 + i])
        if cid > 0 and len(hand_card_ids) < hand_size:
            hand_card_ids.append(cid)

    # Board creatures (indices [15-64], 5 slots x 10 floats)
    board: list[dict[str, Any]] = []
    for slot in range(CREATURE_SLOTS):
        base = offset + 15 + slot * 10
        occupied = obs[base + 0] > 0.5
        if occupied:
            creature: dict[str, Any] = {
                "slot": slot,
                "card_id": 0,  # Filled in later from embeddings
                "attack": round(obs[base + 1] * 10),
                "health": round(obs[base + 2] * 10),
                "max_health": round(obs[base + 3] * 10),
                "keywords": _decode_keywords(obs[base + 9]),
            }
            board.append(creature)

    return {
        "life": life,
        "essence": essence,
        "ap": ap,
        "deck_size": deck_size,
        "hand_size": hand_size,
        "hand_card_ids": hand_card_ids,
        "board": board,
    }


def _assign_creature_card_ids(
    obs: np.ndarray[Any, Any],
    p1: dict[str, Any],
    p2: dict[str, Any],
) -> None:
    """Assign card IDs to board creatures from the embeddings section.

    The embeddings (indices 156-325) are packed sequentially:
    P1 hand → P1 creatures → P1 supports → P2 hand → P2 creatures → P2 supports
    """
    idx = 156

    # P1 hand (skip, already decoded from per-player section)
    idx += len(p1["hand_card_ids"])

    # P1 creatures
    for creature in p1["board"]:
        if idx < 326:
            creature["card_id"] = _decode_card_id(obs[idx])
            idx += 1

    # P1 supports (skip, count from tensor)
    p1_off = 6
    for slot in range(SUPPORT_SLOTS):
        base = p1_off + 65 + slot * 5
        if obs[base] > 0.5:
            idx += 1

    # P2 hand (skip)
    idx += len(p2["hand_card_ids"])

    # P2 creatures
    for creature in p2["board"]:
        if idx < 326:
            creature["card_id"] = _decode_card_id(obs[idx])
            idx += 1


# ============================================================================
# Action Decoder (crates/cardgame/src/core/actions.rs)
# ============================================================================

def decode_action(
    action_index: int,
    hand_card_ids: list[int] | None = None,
) -> dict[str, Any]:
    """Decode action index (0-421) into structured action dict.

    Args:
        action_index: The action index from the engine.
        hand_card_ids: Optional list of card IDs in the player's hand,
            used to annotate PlayCard with the actual card ID.
    """
    if action_index <= 119:
        # PlayCard: hand_index * 12 + target_slot
        hand_index = action_index // 12
        target_slot = action_index % 12
        card_id = 0
        if hand_card_ids and hand_index < len(hand_card_ids):
            card_id = hand_card_ids[hand_index]
        return {
            "type": "PlayCard",
            "hand_index": hand_index,
            "target_slot": target_slot,
            "card_id": card_id,
        }
    elif action_index <= 144:
        # Attack: 120 + attacker * 5 + defender
        offset = action_index - 120
        attacker = offset // 5
        defender = offset % 5
        return {
            "type": "Attack",
            "attacker_slot": attacker,
            "defender_slot": defender,
        }
    elif action_index <= 419:
        # UseAbility: 145 + slot * 55 + ability * 11 + target
        offset = action_index - 145
        slot = offset // 55
        remainder = offset % 55
        ability = remainder // 11
        target = remainder % 11
        target_str: str | int
        if target == 0:
            target_str = "none"
        elif target <= 5:
            target_str = target - 1  # enemy slot 0-4
        else:
            target_str = f"ally:{target - 6}"  # ally slot 0-4
        return {
            "type": "UseAbility",
            "slot": slot,
            "ability_index": ability,
            "target": target_str,
        }
    elif action_index == 420:
        return {"type": "CommanderInsight"}
    elif action_index == 421:
        return {"type": "EndTurn"}
    else:
        return {"type": "EndTurn"}


def decode_legal_action_indices(
    mask: np.ndarray[Any, Any],
    hand_card_ids: list[int] | None = None,
) -> list[dict[str, Any]]:
    """Decode all legal actions from the action mask."""
    actions = []
    for i in range(422):
        if mask[i] > 0.5:
            action = decode_action(i, hand_card_ids)
            action["_index"] = i
            actions.append(action)
    return actions


# ============================================================================
# Game Recording
# ============================================================================

def _get_bot_action(game: PyGame, bot: str) -> int:
    """Get action from specified bot type."""
    if bot == "greedy":
        return game.greedy_action()
    elif bot == "random":
        return game.random_action()
    elif bot.startswith("mcts"):
        sims = int(bot[4:])
        return game.mcts_action(sims)
    elif bot.startswith("alphabeta"):
        depth = int(bot[9:])
        return game.alphabeta_action(depth)
    else:
        msg = f"Unknown bot: {bot}. Use greedy, random, mcts<N>, alphabeta<N>"
        raise ValueError(msg)


def play_game(
    deck1: str,
    deck2: str,
    seed: int,
    bot: str = "greedy",
    with_action_values: bool = False,
    av_rollouts: int = 1,
    with_reasoning: bool = False,
    reasoning_depth: int = 3,
) -> dict[str, Any]:
    """Play one game and record a structured transcript.

    Args:
        deck1: P1 deck name
        deck2: P2 deck name
        seed: Random seed for reproducibility
        bot: Bot type for playing both sides
        with_action_values: If True, compute per-action values via rollouts
        av_rollouts: Number of rollouts per legal action for AV computation
        with_reasoning: If True, annotate positions with AB-ranked evaluations
            as chain-of-thought reasoning traces
        reasoning_depth: AB search depth for reasoning traces (default: 3)

    Returns:
        Game transcript dict with header, positions, and result.
    """
    game = PyGame(deck1=deck1, deck2=deck2)
    game.reset(seed=seed)

    obs = game.observe()
    commander1_id = _decode_card_id(obs[326])
    commander2_id = _decode_card_id(obs[327])

    header = {
        "deck1": deck1,
        "deck2": deck2,
        "mode": "essence_war",
        "commander1_id": commander1_id,
        "commander2_id": commander2_id,
        "seed": seed,
        "bot": bot,
    }
    if with_reasoning:
        header["reasoning_depth"] = reasoning_depth

    positions: list[dict[str, Any]] = []

    while not game.is_done():
        obs = game.observe()
        mask = game.action_mask()
        state = decode_observation(obs)
        current_player = game.current_player()

        # Get hand cards for the active player
        hand_cards = state["p1_hand"] if current_player == 0 else state["_p2_hand"]

        # Optimization: if bot is AB at same depth as reasoning, do one search
        ab_ranked: list[tuple[int, float]] | None = None
        if with_reasoning and bot.startswith("alphabeta"):
            bot_depth = int(bot[9:])
            if bot_depth == reasoning_depth:
                # Single search serves both purposes
                ab_ranked = game.alphabeta_ranked_actions(depth=reasoning_depth)
                action_idx = ab_ranked[0][0] if ab_ranked else 255
            else:
                action_idx = _get_bot_action(game, bot)
        else:
            action_idx = _get_bot_action(game, bot)

        action = decode_action(action_idx, hand_cards)

        # Legal action indices (compact representation of the 422-mask)
        legal_indices = [int(i) for i in range(422) if mask[i] > 0.5]

        # Build position record
        position: dict[str, Any] = {
            "state": {k: v for k, v in state.items() if not k.startswith("_")},
            "action": action,
            "action_index": action_idx,
            "active_player": current_player,
            "legal_action_indices": legal_indices,
        }

        # AB-ranked reasoning trace (chain-of-thought from search)
        if with_reasoning:
            if ab_ranked is None:
                ab_ranked = game.alphabeta_ranked_actions(depth=reasoning_depth)
            reasoning = _build_reasoning_trace(ab_ranked, hand_cards)
            position["reasoning_trace"] = reasoning

        # Compute action values via greedy rollouts (optional)
        if with_action_values:
            legal_actions = decode_legal_action_indices(mask, hand_cards)
            _annotate_action_values(game, legal_actions, av_rollouts)
            position["legal_actions"] = legal_actions

        positions.append(position)

        # Apply the action
        _reward, _done = game.step(action_idx)

    # Determine game result
    p1_reward = game.get_reward(0)
    if p1_reward > 0:
        winner = 0
    elif p1_reward < 0:
        winner = 1
    else:
        winner = -1  # draw

    result = {
        "winner": winner,
        "reason": "life_zero",  # Most common; engine doesn't expose reason directly
        "total_turns": game.turn_number(),
        "total_positions": len(positions),
    }

    # Annotate positions with game outcome as state value
    for pos in positions:
        player = pos["active_player"]
        if winner == -1:
            pos["state_value"] = 0.0
        elif player == winner:
            pos["state_value"] = 1.0
        else:
            pos["state_value"] = -1.0

    return {
        "header": header,
        "positions": positions,
        "result": result,
    }


def _build_reasoning_trace(
    ranked: list[tuple[int, float]],
    hand_card_ids: list[int] | None = None,
) -> list[dict[str, Any]]:
    """Convert AB-ranked (action_index, score) pairs to structured trace.

    Args:
        ranked: List of (action_index, raw_score) from alphabeta_ranked_actions.
        hand_card_ids: Hand cards for PlayCard annotation.

    Returns:
        List of action dicts with 'raw_score' key, sorted best-first.
    """
    trace = []
    for action_idx, raw_score in ranked:
        action = decode_action(action_idx, hand_card_ids)
        action["raw_score"] = raw_score
        trace.append(action)
    return trace


def _annotate_action_values(
    game: PyGame,
    legal_actions: list[dict[str, Any]],
    num_rollouts: int,
) -> None:
    """Compute action values via greedy rollouts from forked game states.

    For each legal action, fork the game, apply the action, then play
    to completion with the greedy bot. The win/loss outcome is the
    action value estimate.
    """
    current_player = game.current_player()

    for la in legal_actions:
        action_idx = la["_index"]
        total_value = 0.0

        for _ in range(num_rollouts):
            forked = game.fork()
            reward, done = forked.step(action_idx)

            if not done:
                # Play to completion with greedy
                while not forked.is_done():
                    greedy_act = forked.greedy_action()
                    forked.step(greedy_act)
                reward = forked.get_reward(current_player)

            total_value += reward

        la["value"] = total_value / num_rollouts


# ============================================================================
# Bulk Generation
# ============================================================================

def _worker_generate_chunk(args: tuple[Any, ...]) -> list[str]:
    """Worker function: generate a chunk of games, return JSONL strings.

    Each worker is an independent process with its own Rust engine instances.
    """
    chunk_matchups, base_seed, chunk_offset, bot, with_av, av_rollouts, with_reasoning, reasoning_depth = args

    results: list[str] = []
    for i, (deck1, deck2) in enumerate(chunk_matchups):
        seed = base_seed + chunk_offset + i
        transcript = play_game(
            deck1=deck1,
            deck2=deck2,
            seed=seed,
            bot=bot,
            with_action_values=with_av,
            av_rollouts=av_rollouts,
            with_reasoning=with_reasoning,
            reasoning_depth=reasoning_depth,
        )
        results.append(json.dumps(transcript, separators=(",", ":")))
    return results


def generate_dataset(
    num_games: int,
    output_path: str | Path,
    bot: str = "greedy",
    with_action_values: bool = False,
    av_rollouts: int = 1,
    matchup_mode: str = "all",
    base_seed: int = 42,
    with_reasoning: bool = False,
    reasoning_depth: int = 3,
    num_workers: int = 0,
) -> dict[str, Any]:
    """Generate a dataset of game transcripts (parallel or sequential).

    Args:
        num_games: Total number of games to generate.
        output_path: Path to output JSONL or JSONL.GZ file.
        bot: Bot type for both players.
        with_action_values: Compute per-action values.
        av_rollouts: Rollouts per action for AV computation.
        matchup_mode: "all" for round-robin, "random" for random pairs.
        base_seed: Starting seed for reproducibility.
        with_reasoning: Annotate positions with AB-ranked reasoning traces.
        reasoning_depth: AB search depth for reasoning traces.
        num_workers: Parallel workers (0 = auto = CPU count - 1).

    Returns:
        Summary statistics dict.
    """
    if num_workers == 0:
        num_workers = max(1, mp.cpu_count() - 1)
    # For small jobs or single-core, fall back to sequential
    if num_workers <= 1 or num_games < 100:
        return _generate_dataset_sequential(
            num_games=num_games,
            output_path=output_path,
            bot=bot,
            with_action_values=with_action_values,
            av_rollouts=av_rollouts,
            matchup_mode=matchup_mode,
            base_seed=base_seed,
            with_reasoning=with_reasoning,
            reasoning_depth=reasoning_depth,
        )

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    use_gzip = str(output_path).endswith(".gz")

    matchups = _generate_matchups(num_games, matchup_mode)

    # Split matchups into chunks for parallel workers
    chunk_size = max(1, num_games // (num_workers * 4))  # ~4 chunks per worker for balancing
    chunks = []
    for start in range(0, num_games, chunk_size):
        end = min(start + chunk_size, num_games)
        chunk = matchups[start:end]
        chunks.append((
            chunk, base_seed, start, bot,
            with_action_values, av_rollouts,
            with_reasoning, reasoning_depth,
        ))

    print(f"Generating {num_games:,} games with {num_workers} workers ({len(chunks)} chunks)")

    start_time = time.time()
    total_positions = 0
    games_by_winner: dict[int, int] = {0: 0, 1: 0, -1: 0}
    turn_counts: list[int] = []
    position_counts: list[int] = []
    games_done = 0

    opener = gzip.open if use_gzip else open
    with opener(output_path, "wt") as f, mp.Pool(num_workers) as pool:
        for chunk_results in pool.imap_unordered(_worker_generate_chunk, chunks):
            for line in chunk_results:
                f.write(line + "\n")
                transcript = json.loads(line)
                n_pos = transcript["result"]["total_positions"]
                total_positions += n_pos
                position_counts.append(n_pos)
                turn_counts.append(transcript["result"]["total_turns"])
                games_by_winner[transcript["result"]["winner"]] += 1
                games_done += 1

            # Progress per chunk
            elapsed = time.time() - start_time
            gps = games_done / elapsed if elapsed > 0 else 0
            eta = (num_games - games_done) / gps if gps > 0 else 0
            print(
                f"  [{games_done:>{len(str(num_games))}}/{num_games}] "
                f"{gps:.0f} games/sec | "
                f"{total_positions:,} positions | "
                f"ETA: {eta:.0f}s"
            )

    elapsed = time.time() - start_time
    total_tokens_est = total_positions * 60

    stats = {
        "total_games": num_games,
        "total_positions": total_positions,
        "estimated_tokens": total_tokens_est,
        "avg_positions_per_game": round(total_positions / num_games, 1),
        "avg_turns_per_game": round(sum(turn_counts) / len(turn_counts), 1),
        "max_positions": max(position_counts),
        "min_positions": min(position_counts),
        "p1_wins": games_by_winner[0],
        "p2_wins": games_by_winner[1],
        "draws": games_by_winner[-1],
        "elapsed_seconds": round(elapsed, 1),
        "games_per_second": round(num_games / elapsed, 1) if elapsed > 0 else 0,
        "bot": bot,
        "with_action_values": with_action_values,
        "with_reasoning": with_reasoning,
        "reasoning_depth": reasoning_depth if with_reasoning else None,
        "num_workers": num_workers,
        "output_path": str(output_path),
    }

    return stats


def _generate_dataset_sequential(
    num_games: int,
    output_path: str | Path,
    bot: str = "greedy",
    with_action_values: bool = False,
    av_rollouts: int = 1,
    matchup_mode: str = "all",
    base_seed: int = 42,
    with_reasoning: bool = False,
    reasoning_depth: int = 3,
) -> dict[str, Any]:
    """Sequential (single-process) dataset generation. Fallback for small jobs."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    use_gzip = str(output_path).endswith(".gz")

    matchups = _generate_matchups(num_games, matchup_mode)

    total_positions = 0
    games_by_winner: dict[int, int] = {0: 0, 1: 0, -1: 0}
    turn_counts: list[int] = []
    position_counts: list[int] = []

    start_time = time.time()
    opener = gzip.open if use_gzip else open

    with opener(output_path, "wt") as f:
        for i, (deck1, deck2) in enumerate(matchups):
            seed = base_seed + i
            transcript = play_game(
                deck1=deck1, deck2=deck2, seed=seed, bot=bot,
                with_action_values=with_action_values, av_rollouts=av_rollouts,
                with_reasoning=with_reasoning, reasoning_depth=reasoning_depth,
            )
            f.write(json.dumps(transcript, separators=(",", ":")) + "\n")

            n_pos = transcript["result"]["total_positions"]
            total_positions += n_pos
            position_counts.append(n_pos)
            turn_counts.append(transcript["result"]["total_turns"])
            games_by_winner[transcript["result"]["winner"]] += 1

            if (i + 1) % max(1, num_games // 20) == 0 or i == num_games - 1:
                elapsed = time.time() - start_time
                gps = (i + 1) / elapsed if elapsed > 0 else 0
                eta = (num_games - i - 1) / gps if gps > 0 else 0
                print(
                    f"  [{i + 1:>{len(str(num_games))}}/{num_games}] "
                    f"{gps:.0f} games/sec | "
                    f"{total_positions:,} positions | "
                    f"ETA: {eta:.0f}s"
                )

    elapsed = time.time() - start_time
    return {
        "total_games": num_games,
        "total_positions": total_positions,
        "estimated_tokens": total_positions * 60,
        "avg_positions_per_game": round(total_positions / num_games, 1),
        "avg_turns_per_game": round(sum(turn_counts) / len(turn_counts), 1),
        "max_positions": max(position_counts),
        "min_positions": min(position_counts),
        "p1_wins": games_by_winner[0],
        "p2_wins": games_by_winner[1],
        "draws": games_by_winner[-1],
        "elapsed_seconds": round(elapsed, 1),
        "games_per_second": round(num_games / elapsed, 1) if elapsed > 0 else 0,
        "bot": bot,
        "with_action_values": with_action_values,
        "with_reasoning": with_reasoning,
        "reasoning_depth": reasoning_depth if with_reasoning else None,
        "num_workers": 1,
        "output_path": str(output_path),
    }


def _generate_matchups(
    num_games: int,
    mode: str,
) -> list[tuple[str, str]]:
    """Generate deck matchup pairs."""
    if mode == "all":
        # Round-robin across all 144 matchups (12 x 12)
        all_pairs = [(d1, d2) for d1 in ALL_DECKS for d2 in ALL_DECKS]
        matchups: list[tuple[str, str]] = []
        while len(matchups) < num_games:
            matchups.extend(all_pairs)
        return matchups[:num_games]
    elif mode == "random":
        rng = random.Random(42)
        return [(rng.choice(ALL_DECKS), rng.choice(ALL_DECKS)) for _ in range(num_games)]
    elif mode == "mirror":
        # Mirror matches only (same deck vs same deck)
        matchups = []
        while len(matchups) < num_games:
            for deck in ALL_DECKS:
                matchups.append((deck, deck))
        return matchups[:num_games]
    else:
        msg = f"Unknown matchup mode: {mode}. Use 'all', 'random', or 'mirror'"
        raise ValueError(msg)


# ============================================================================
# CLI
# ============================================================================

def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate game transcripts for Game Transformer training"
    )
    parser.add_argument(
        "--games", type=int, default=1000,
        help="Number of games to generate (default: 1000)",
    )
    parser.add_argument(
        "--bot", type=str, default="greedy",
        help="Bot type: greedy, random, mcts<N>, alphabeta<N> (default: greedy)",
    )
    parser.add_argument(
        "--output", type=str, default="data/transcripts/games.jsonl.gz",
        help="Output file path (supports .jsonl and .jsonl.gz)",
    )
    parser.add_argument(
        "--with-av", action="store_true",
        help="Compute per-action values via greedy rollouts",
    )
    parser.add_argument(
        "--av-rollouts", type=int, default=1,
        help="Rollouts per action for AV computation (default: 1)",
    )
    parser.add_argument(
        "--matchup-mode", type=str, default="all",
        choices=["all", "random", "mirror"],
        help="Matchup selection: all (round-robin), random, mirror (default: all)",
    )
    parser.add_argument(
        "--seed", type=int, default=42,
        help="Base random seed (default: 42)",
    )
    parser.add_argument(
        "--with-reasoning", action="store_true",
        help="Annotate positions with AB search reasoning traces (chain-of-thought)",
    )
    parser.add_argument(
        "--reasoning-depth", type=int, default=3,
        help="Alpha-Beta search depth for reasoning traces (default: 3)",
    )
    parser.add_argument(
        "--workers", type=int, default=0,
        help="Parallel workers (0 = auto = CPU count - 1)",
    )
    args = parser.parse_args()

    print(f"Generating {args.games:,} games with {args.bot} bot")
    print(f"  Output: {args.output}")
    print(f"  Matchups: {args.matchup_mode}")
    print(f"  Action values: {'yes' if args.with_av else 'no'}")
    if args.with_av:
        print(f"  AV rollouts: {args.av_rollouts}")
    if args.with_reasoning:
        print(f"  Reasoning traces: AB depth {args.reasoning_depth}")
    print()

    stats = generate_dataset(
        num_games=args.games,
        output_path=args.output,
        bot=args.bot,
        with_action_values=args.with_av,
        av_rollouts=args.av_rollouts,
        matchup_mode=args.matchup_mode,
        base_seed=args.seed,
        with_reasoning=args.with_reasoning,
        reasoning_depth=args.reasoning_depth,
        num_workers=args.workers,
    )

    print("\nGeneration complete!")
    print(f"  Games: {stats['total_games']:,}")
    print(f"  Positions: {stats['total_positions']:,}")
    print(f"  Est. tokens: {stats['estimated_tokens']:,}")
    print(f"  Avg positions/game: {stats['avg_positions_per_game']}")
    print(f"  Avg turns/game: {stats['avg_turns_per_game']}")
    print(f"  P1 wins: {stats['p1_wins']:,} | P2 wins: {stats['p2_wins']:,} | Draws: {stats['draws']:,}")
    print(f"  Speed: {stats['games_per_second']:.0f} games/sec")
    print(f"  Time: {stats['elapsed_seconds']:.1f}s")
    print(f"  Output: {stats['output_path']}")

    # Save stats alongside output
    stats_path = Path(args.output).with_suffix(".stats.json")
    with stats_path.open("w") as f:
        json.dump(stats, f, indent=2)
    print(f"  Stats: {stats_path}")


if __name__ == "__main__":
    main()
