"""DAgger (Dataset Aggregation) data generation for Game Transformer.

Lets the trained model play games (encountering its own state distribution),
then annotates every position with expert reasoning (Alpha-Beta search).
The resulting transcripts are in the same JSONL format as data_generator.py,
so existing preprocessing and training pipelines work unchanged.

Two-phase pipeline for efficiency:
  Phase 1 (GPU): Batched model inference — maintain pool of concurrent games,
    single forward pass per step. Records (deck1, deck2, seed, action_indices).
  Phase 2 (CPU): Parallel annotation — replay games from recorded actions,
    call AB search at each position for reasoning traces. Workers don't need
    the model, so this is embarrassingly parallel.

Usage:
    uv run python -m essence_wars.transformer.dagger \
        --checkpoint experiments/transformer/ew-small_greedy_100k_reasoning_e10/best.pt \
        --games 1000 --output data/transcripts/dagger_iter0_1k.jsonl.gz
"""

from __future__ import annotations

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

import numpy as np
import torch

from essence_wars import PyGame
from essence_wars.transformer.data_generator import (
    _build_reasoning_trace,
    _decode_card_id,
    _generate_matchups,
    decode_action,
    decode_observation,
)
from essence_wars.transformer.model import GameTransformer
from essence_wars.transformer.tokenizer import EssenceWarsTokenizer

# ============================================================================
# Tokenization helper (also used by phase6_dagger.py evaluation)
# ============================================================================


def _tokenize_for_model(
    tok: EssenceWarsTokenizer,
    max_seq_len: int,
    state: dict[str, Any],
    game_header: dict[str, Any],
) -> tuple[list[int], list[int]]:
    """Encode game_header + state into padded token IDs and attention mask.

    Returns (input_ids, attention_mask) as plain lists, ready for tensor conversion.
    """
    tokens = tok.encode_game_header(game_header)
    tokens.append("[SEP]")
    tokens.extend(tok.encode_state(state))
    context_ids = tok.encode_tokens(tokens)

    seq_len = min(len(context_ids), max_seq_len)
    pad_len = max_seq_len - seq_len

    input_ids = context_ids[:seq_len] + [tok.pad_token_id] * pad_len
    attn_mask = [1] * seq_len + [0] * pad_len

    return input_ids, attn_mask


# ============================================================================
# Phase 1: GPU batched game play
# ============================================================================


def _play_games_on_gpu(
    checkpoint_path: Path,
    matchups: list[tuple[str, str]],
    max_seq_len: int,
    base_seed: int,
    pool_size: int = 32,
) -> list[dict[str, Any]]:
    """Play games using batched GPU inference. Returns lightweight game records.

    Maintains a pool of concurrent games. Each step:
    1. Tokenize all active game states → single batched forward pass on GPU
    2. Select greedy actions from model output
    3. Step all games, replace finished ones from queue

    Returns list of game records: {deck1, deck2, seed, action_indices, commander1_id,
    commander2_id, winner, total_turns}. These are replayed in Phase 2 for annotation.
    """
    # Load model on GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    config = checkpoint["config"]
    model = GameTransformer(config)
    model.load_state_dict(checkpoint["model_state_dict"])
    model = model.to(device)
    model.eval()

    tok = EssenceWarsTokenizer()
    num_games = len(matchups)
    pool_size = min(pool_size, num_games)

    # Active game slots
    games: list[PyGame | None] = [None] * pool_size
    game_headers: list[dict[str, Any]] = [{}] * pool_size
    game_records: list[dict[str, Any]] = [{}] * pool_size
    next_game_idx = 0  # Next matchup to start

    completed: list[dict[str, Any]] = []
    start_time = time.time()

    def _init_game(slot: int) -> bool:
        """Initialize a game in the given slot. Returns False if no games left."""
        nonlocal next_game_idx
        if next_game_idx >= num_games:
            return False
        deck1, deck2 = matchups[next_game_idx]
        seed = base_seed + next_game_idx

        game = PyGame(deck1=deck1, deck2=deck2)
        game.reset(seed=seed)

        obs = game.observe()
        c1_id = _decode_card_id(obs[326])
        c2_id = _decode_card_id(obs[327])

        games[slot] = game
        game_headers[slot] = {
            "deck1": deck1,
            "deck2": deck2,
            "mode": "essence_war",
            "commander1_id": c1_id,
            "commander2_id": c2_id,
        }
        game_records[slot] = {
            "deck1": deck1,
            "deck2": deck2,
            "seed": seed,
            "commander1_id": c1_id,
            "commander2_id": c2_id,
            "action_indices": [],
        }
        next_game_idx += 1
        return True

    # Fill initial pool
    for slot in range(pool_size):
        _init_game(slot)

    # Main play loop
    total_steps = 0
    while True:
        # Find active (non-done) games
        active_slots = []
        for s in range(pool_size):
            g = games[s]
            if g is not None and not g.is_done():
                active_slots.append(s)
        if not active_slots:
            break

        # Tokenize all active games
        all_input_ids = []
        all_attn_masks = []
        all_legal_masks = []

        for s in active_slots:
            game = games[s]
            assert game is not None  # guaranteed by active_slots filter
            obs = game.observe()
            mask = game.action_mask()
            state = decode_observation(obs)
            clean_state = {k: v for k, v in state.items() if not k.startswith("_")}

            input_ids, attn_mask = _tokenize_for_model(
                tok, max_seq_len, clean_state, game_headers[s]
            )
            all_input_ids.append(input_ids)
            all_attn_masks.append(attn_mask)
            all_legal_masks.append(mask)

        # Batched forward pass
        input_t = torch.tensor(all_input_ids, dtype=torch.long, device=device)
        attn_t = torch.tensor(all_attn_masks, dtype=torch.long, device=device)
        legal_t = torch.from_numpy(np.array(all_legal_masks)).float().to(device)

        action_indices, _ = model.select_action(input_t, attn_t, legal_t, mode="greedy")

        # Step all games
        for i, s in enumerate(active_slots):
            action_idx = int(action_indices[i].item())
            game_records[s]["action_indices"].append(action_idx)
            g = games[s]
            assert g is not None
            g.step(action_idx)
            total_steps += 1

        # Check for finished games, replace with new ones
        for s in active_slots:
            g = games[s]
            assert g is not None
            if g.is_done():
                p1_reward = g.get_reward(0)
                if p1_reward > 0:
                    winner = 0
                elif p1_reward < 0:
                    winner = 1
                else:
                    winner = -1

                record = game_records[s]
                record["winner"] = winner
                record["total_turns"] = g.turn_number()
                completed.append(record)

                # Progress reporting
                if len(completed) % max(1, num_games // 20) == 0 or len(completed) == num_games:
                    elapsed = time.time() - start_time
                    gps = len(completed) / elapsed if elapsed > 0 else 0
                    eta = (num_games - len(completed)) / gps if gps > 0 else 0
                    print(
                        f"  [Play] {len(completed):>{len(str(num_games))}}/{num_games} "
                        f"| {gps:.0f} games/sec | {total_steps:,} steps | ETA: {eta:.0f}s"
                    )

                # Replace with new game
                if not _init_game(s):
                    games[s] = None

    elapsed = time.time() - start_time
    print(
        f"  [Play] Done: {len(completed):,} games, {total_steps:,} steps "
        f"in {elapsed:.0f}s ({len(completed) / elapsed:.0f} games/sec)"
    )

    # Free GPU memory before annotation phase
    del model, checkpoint
    torch.cuda.empty_cache()

    return completed


# ============================================================================
# Phase 2: CPU parallel annotation
# ============================================================================


def _annotate_worker(args: tuple[Any, ...]) -> list[str]:
    """Worker: replay games from recorded actions, annotate with AB reasoning.

    Each worker replays a chunk of games. At each position, it calls
    alphabeta_ranked_actions() for the expert reasoning trace. Workers
    don't need the model — just the game engine.
    """
    game_records, reasoning_depth = args

    results: list[str] = []
    for record in game_records:
        transcript = _replay_and_annotate(record, reasoning_depth)
        results.append(json.dumps(transcript, separators=(",", ":")))
    return results


def _replay_and_annotate(
    record: dict[str, Any],
    reasoning_depth: int,
) -> dict[str, Any]:
    """Replay a single game from recorded actions and annotate each position.

    Args:
        record: Game record from Phase 1 with deck1, deck2, seed, action_indices,
            commander1_id, commander2_id, winner, total_turns.
        reasoning_depth: AB search depth for expert annotations.

    Returns:
        Full transcript dict in data_generator.play_game() format.
    """
    deck1 = record["deck1"]
    deck2 = record["deck2"]
    seed = record["seed"]
    action_indices = record["action_indices"]

    game = PyGame(deck1=deck1, deck2=deck2)
    game.reset(seed=seed)

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

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

    for action_idx in action_indices:
        obs = game.observe()
        mask = game.action_mask()
        state = decode_observation(obs)
        current_player = game.current_player()

        hand_cards = state["p1_hand"] if current_player == 0 else state["_p2_hand"]
        clean_state = {k: v for k, v in state.items() if not k.startswith("_")}
        action = decode_action(action_idx, hand_cards)

        # Expert annotation: AB-ranked reasoning trace
        ab_ranked = game.alphabeta_ranked_actions(depth=reasoning_depth)
        reasoning = _build_reasoning_trace(ab_ranked, hand_cards)

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

        position: dict[str, Any] = {
            "state": clean_state,
            "action": action,
            "action_index": action_idx,
            "active_player": current_player,
            "legal_action_indices": legal_indices,
            "reasoning_trace": reasoning,
        }
        positions.append(position)

        game.step(action_idx)

    # Annotate positions with game outcome as state value
    winner = record["winner"]
    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

    result = {
        "winner": winner,
        "reason": "life_zero",
        "total_turns": record["total_turns"],
        "total_positions": len(positions),
    }

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


def _annotate_transcripts_parallel(
    game_records: list[dict[str, Any]],
    reasoning_depth: int,
    num_workers: int,
    output_path: Path,
) -> dict[str, Any]:
    """Annotate game records in parallel using CPU workers.

    Args:
        game_records: List of game records from Phase 1.
        reasoning_depth: AB depth for expert annotations.
        num_workers: Number of parallel CPU workers.
        output_path: Output file path (.jsonl or .jsonl.gz).

    Returns:
        Stats dict with position counts and game outcomes.
    """
    use_gzip = str(output_path).endswith(".gz")
    num_records = len(game_records)

    # Split into chunks for workers
    chunk_size = max(1, num_records // (num_workers * 4))
    chunks = []
    for start in range(0, num_records, chunk_size):
        end = min(start + chunk_size, num_records)
        chunks.append((game_records[start:end], reasoning_depth))

    print(
        f"  [Annotate] {num_records:,} games with {num_workers} workers "
        f"({len(chunks)} chunks, AB-d{reasoning_depth})"
    )

    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

    # Write to plain JSONL first, then compress — avoids gzip stream corruption
    # during long write sessions with large data volumes.
    tmp_path = output_path.with_suffix(".jsonl") if use_gzip else output_path
    with tmp_path.open("w") as f, mp.Pool(num_workers) as pool:
        for chunk_results in pool.imap_unordered(_annotate_worker, 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

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

    # Compress after all writes are complete (ensures gzip integrity)
    if use_gzip:
        print("  [Annotate] Compressing output...")
        with tmp_path.open("rb") as f_in, gzip.open(output_path, "wb") as f_out:
            while chunk := f_in.read(8 * 1024 * 1024):
                f_out.write(chunk)
        tmp_path.unlink()

    elapsed = time.time() - start_time
    print(
        f"  [Annotate] Done: {games_done:,} games, {total_positions:,} positions "
        f"in {elapsed:.0f}s ({games_done / elapsed:.0f} games/sec)"
    )

    return {
        "total_positions": total_positions,
        "avg_positions_per_game": round(total_positions / num_records, 1) if num_records else 0,
        "avg_turns_per_game": round(sum(turn_counts) / len(turn_counts), 1) if turn_counts else 0,
        "max_positions": max(position_counts) if position_counts else 0,
        "min_positions": min(position_counts) if position_counts else 0,
        "p1_wins": games_by_winner[0],
        "p2_wins": games_by_winner[1],
        "draws": games_by_winner[-1],
        "annotate_seconds": round(elapsed, 1),
    }


# ============================================================================
# Public API
# ============================================================================


def generate_dagger_dataset(
    checkpoint_path: str | Path,
    num_games: int,
    output_path: str | Path,
    reasoning_depth: int = 3,
    max_seq_len: int = 192,
    matchup_mode: str = "all",
    base_seed: int = 1_000_000,
    num_workers: int = 0,
    pool_size: int = 32,
) -> dict[str, Any]:
    """Generate a DAgger dataset: model plays on GPU, expert annotates on CPU.

    Two-phase pipeline:
      Phase 1: Batched GPU inference — model plays all games with concurrent
        game pool. Fast (~100+ games/sec on RTX 4090).
      Phase 2: CPU annotation — replay each game and annotate with AB search.
        Embarrassingly parallel across games.

    Args:
        checkpoint_path: Path to model checkpoint (.pt file).
        num_games: Total number of games to generate.
        output_path: Output JSONL or JSONL.GZ file path.
        reasoning_depth: AB depth for expert annotations.
        max_seq_len: Token sequence length for model inference.
        matchup_mode: Matchup selection mode (all/random/mirror).
        base_seed: Starting seed for game generation.
        num_workers: CPU workers for annotation (0 = auto).
        pool_size: Concurrent games during GPU play phase.

    Returns:
        Summary statistics dict.
    """
    checkpoint_path = Path(checkpoint_path)
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if num_workers == 0:
        num_workers = max(1, mp.cpu_count() - 1)

    matchups = _generate_matchups(num_games, matchup_mode)
    total_start = time.time()

    # Phase 1: GPU play
    print(f"\n  Phase 1: Playing {num_games:,} games (GPU, pool={pool_size})")
    game_records = _play_games_on_gpu(
        checkpoint_path=checkpoint_path,
        matchups=matchups,
        max_seq_len=max_seq_len,
        base_seed=base_seed,
        pool_size=pool_size,
    )

    # Phase 2: CPU annotation
    print(f"\n  Phase 2: Annotating with AB-d{reasoning_depth} ({num_workers} workers)")
    annotation_stats = _annotate_transcripts_parallel(
        game_records=game_records,
        reasoning_depth=reasoning_depth,
        num_workers=num_workers,
        output_path=output_path,
    )

    total_elapsed = time.time() - total_start
    stats: dict[str, Any] = {
        "total_games": num_games,
        "total_positions": annotation_stats["total_positions"],
        "estimated_tokens": annotation_stats["total_positions"] * 60,
        "avg_positions_per_game": annotation_stats["avg_positions_per_game"],
        "avg_turns_per_game": annotation_stats["avg_turns_per_game"],
        "max_positions": annotation_stats["max_positions"],
        "min_positions": annotation_stats["min_positions"],
        "p1_wins": annotation_stats["p1_wins"],
        "p2_wins": annotation_stats["p2_wins"],
        "draws": annotation_stats["draws"],
        "elapsed_seconds": round(total_elapsed, 1),
        "games_per_second": round(num_games / total_elapsed, 1) if total_elapsed > 0 else 0,
        "annotate_seconds": annotation_stats["annotate_seconds"],
        "method": "dagger_gpu",
        "reasoning_depth": reasoning_depth,
        "checkpoint": str(checkpoint_path),
        "num_workers": num_workers,
        "pool_size": pool_size,
        "output_path": str(output_path),
    }
    return stats


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


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate DAgger dataset: model plays (GPU), expert annotates (CPU)"
    )
    parser.add_argument(
        "--checkpoint", type=str, required=True,
        help="Path to model checkpoint (.pt file)",
    )
    parser.add_argument(
        "--games", type=int, default=1000,
        help="Number of games to generate (default: 1000)",
    )
    parser.add_argument(
        "--output", type=str, default="data/transcripts/dagger.jsonl.gz",
        help="Output file path (supports .jsonl and .jsonl.gz)",
    )
    parser.add_argument(
        "--reasoning-depth", type=int, default=3,
        help="AB search depth for expert annotations (default: 3)",
    )
    parser.add_argument(
        "--max-seq-len", type=int, default=192,
        help="Token sequence length for model inference (default: 192)",
    )
    parser.add_argument(
        "--matchup-mode", type=str, default="all",
        choices=["all", "random", "mirror"],
        help="Matchup selection mode (default: all)",
    )
    parser.add_argument(
        "--seed", type=int, default=1_000_000,
        help="Base random seed (default: 1000000)",
    )
    parser.add_argument(
        "--workers", type=int, default=0,
        help="CPU workers for annotation (0 = auto = CPU count - 1)",
    )
    parser.add_argument(
        "--pool-size", type=int, default=32,
        help="Concurrent games during GPU play phase (default: 32)",
    )
    args = parser.parse_args()

    print(f"DAgger generation: {args.games:,} games")
    print(f"  Checkpoint: {args.checkpoint}")
    print(f"  Expert depth: AB-d{args.reasoning_depth}")
    print(f"  GPU pool size: {args.pool_size}")
    print(f"  Output: {args.output}")
    print()

    stats = generate_dagger_dataset(
        checkpoint_path=args.checkpoint,
        num_games=args.games,
        output_path=args.output,
        reasoning_depth=args.reasoning_depth,
        max_seq_len=args.max_seq_len,
        matchup_mode=args.matchup_mode,
        base_seed=args.seed,
        num_workers=args.workers,
        pool_size=args.pool_size,
    )

    print("\nDAgger generation complete!")
    print(f"  Games: {stats['total_games']:,}")
    print(f"  Positions: {stats['total_positions']:,}")
    print(f"  Avg positions/game: {stats['avg_positions_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()
