"""PyTorch dataset for Game Transformer training.

Two-phase pipeline to avoid OOM:
1. `preprocess()`: Reads JSONL transcripts, tokenizes, saves as compact
   memory-mapped numpy arrays on disk. Processes one game at a time.
2. `MemmappedPositionDataset`: Memory-maps the preprocessed arrays.
   Near-zero RAM usage — only accessed batch pages are loaded by the OS.

Memory comparison (7M positions):
- Old (in-memory dicts): ~6-8 GB RAM → OOM crash
- New (memory-mapped):   ~160 KB per batch → stable on any machine
"""

from __future__ import annotations

import gzip
import json
import random
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.utils.data import Dataset

from essence_wars.transformer.tokenizer import EssenceWarsTokenizer

ACTION_SPACE_SIZE = 422


# ============================================================================
# Preprocessing: JSONL → memory-mapped numpy arrays
# ============================================================================

def preprocess(
    data_paths: str | Path | list[str | Path],
    output_dir: str | Path,
    max_seq_len: int = 128,
    val_fraction: float = 0.1,
    split_seed: int = 42,
    tokenizer: EssenceWarsTokenizer | None = None,
) -> dict[str, Any]:
    """Preprocess JSONL transcripts into memory-mapped numpy arrays.

    Processes one game at a time to keep memory usage minimal.

    Args:
        data_paths: JSONL/JSONL.GZ transcript file(s).
        output_dir: Directory for preprocessed output.
        max_seq_len: Token sequence length (pad/truncate).
        val_fraction: Fraction of games for validation.
        split_seed: Seed for train/val split.
        tokenizer: Tokenizer instance (created if None).

    Returns:
        Stats dict with counts and paths.
    """
    if isinstance(data_paths, (str, Path)):
        data_paths = [data_paths]
    data_paths = [Path(p) for p in data_paths]

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    tok = tokenizer or EssenceWarsTokenizer()

    # --- Pass 1: Count games and positions ---
    print("  Pass 1: Counting positions...")
    t0 = time.time()
    game_position_counts: list[int] = []  # positions per game

    for path in data_paths:
        opener = gzip.open if str(path).endswith(".gz") else open
        with opener(path, "rt") as f:
            for line in f:
                game = json.loads(line)
                game_position_counts.append(len(game["positions"]))

    total_games = len(game_position_counts)
    total_positions = sum(game_position_counts)
    print(f"    {total_games:,} games, {total_positions:,} positions ({time.time()-t0:.1f}s)")

    # --- Train/val split at game level ---
    rng = random.Random(split_seed)
    game_indices = list(range(total_games))
    rng.shuffle(game_indices)
    val_count = max(1, int(total_games * val_fraction))
    val_game_set = set(game_indices[:val_count])

    train_count = sum(
        c for i, c in enumerate(game_position_counts) if i not in val_game_set
    )
    val_count_pos = total_positions - train_count

    print(f"    Train: {train_count:,} positions, Val: {val_count_pos:,} positions")

    # --- Pass 2: Allocate memmaps and fill ---
    print("  Pass 2: Tokenizing and writing arrays...")
    t0 = time.time()

    # Create memory-mapped arrays
    arrays = {
        "train": _create_memmaps(output_dir / "train", train_count, max_seq_len),
        "val": _create_memmaps(output_dir / "val", val_count_pos, max_seq_len),
    }
    train_idx = 0
    val_idx = 0
    game_counter = 0
    processed = 0

    for path in data_paths:
        opener = gzip.open if str(path).endswith(".gz") else open
        with opener(path, "rt") as f:
            for line in f:
                game = json.loads(line)
                is_val = game_counter in val_game_set
                split = "val" if is_val else "train"
                idx_ref = val_idx if is_val else train_idx

                header = game["header"]

                for pos in game["positions"]:
                    _tokenize_and_write(
                        tok, header, pos, arrays[split], idx_ref, max_seq_len
                    )
                    idx_ref += 1
                    processed += 1

                if is_val:
                    val_idx = idx_ref
                else:
                    train_idx = idx_ref

                game_counter += 1

                if game_counter % 10000 == 0:
                    elapsed = time.time() - t0
                    rate = processed / elapsed if elapsed > 0 else 0
                    print(
                        f"    {game_counter:,}/{total_games:,} games "
                        f"({processed:,} positions, {rate:,.0f} pos/sec)"
                    )

    # Flush memmaps
    for split_arrays in arrays.values():
        for arr in split_arrays.values():
            arr.flush()

    elapsed = time.time() - t0
    print(f"    Done! {processed:,} positions in {elapsed:.1f}s ({processed/elapsed:,.0f} pos/sec)")

    # Save metadata
    meta = {
        "total_games": total_games,
        "total_positions": total_positions,
        "train_positions": train_count,
        "val_positions": val_count_pos,
        "max_seq_len": max_seq_len,
        "vocab_size": tok.vocab_size,
        "val_fraction": val_fraction,
        "split_seed": split_seed,
        "data_paths": [str(p) for p in data_paths],
    }
    meta_path = output_dir / "metadata.json"
    with meta_path.open("w") as f:
        json.dump(meta, f, indent=2)

    print(f"  Saved to {output_dir}")
    return meta


def _create_memmaps(
    split_dir: Path, count: int, max_seq_len: int
) -> dict[str, np.memmap[Any, Any]]:
    """Create memory-mapped arrays for one split."""
    split_dir.mkdir(parents=True, exist_ok=True)

    return {
        "input_ids": np.memmap(
            split_dir / "input_ids.dat",
            dtype=np.int16, mode="w+", shape=(count, max_seq_len),
        ),
        "attention_mask": np.memmap(
            split_dir / "attention_mask.dat",
            dtype=np.uint8, mode="w+", shape=(count, max_seq_len),
        ),
        "action_targets": np.memmap(
            split_dir / "action_targets.dat",
            dtype=np.int16, mode="w+", shape=(count,),
        ),
        "value_targets": np.memmap(
            split_dir / "value_targets.dat",
            dtype=np.float32, mode="w+", shape=(count,),
        ),
        "legal_masks": np.memmap(
            split_dir / "legal_masks.dat",
            dtype=np.uint8, mode="w+", shape=(count, ACTION_SPACE_SIZE),
        ),
    }


def _tokenize_and_write(
    tok: EssenceWarsTokenizer,
    header: dict[str, Any],
    pos: dict[str, Any],
    arrays: dict[str, np.memmap[Any, Any]],
    idx: int,
    max_seq_len: int,
) -> None:
    """Tokenize one position and write to memory-mapped arrays."""
    # Build position dict for tokenizer
    tok_input: dict[str, Any] = {
        "game_header": header,
        "state": pos["state"],
        "action": pos["action"],
        "state_value": pos["state_value"],
    }
    if "reasoning_trace" in pos:
        tok_input["reasoning_trace"] = pos["reasoning_trace"]

    # Tokenize context (everything before [ACTION])
    tokens = tok.encode_position(tok_input)
    try:
        action_split = tokens.index("[ACTION]")
    except ValueError:
        action_split = len(tokens)
    context_tokens = tokens[:action_split]
    input_ids = tok.encode_tokens(context_tokens)

    # Pad/truncate
    seq_len = min(len(input_ids), max_seq_len)
    max_seq_len - seq_len

    ids = np.zeros(max_seq_len, dtype=np.int16)
    ids[:seq_len] = input_ids[:seq_len]
    # pad token id = 0 (already zero-initialized)

    mask = np.zeros(max_seq_len, dtype=np.uint8)
    mask[:seq_len] = 1

    # Legal mask
    legal = np.zeros(ACTION_SPACE_SIZE, dtype=np.uint8)
    for i in pos.get("legal_action_indices", []):
        legal[i] = 1

    # Write
    arrays["input_ids"][idx] = ids
    arrays["attention_mask"][idx] = mask
    arrays["action_targets"][idx] = pos["action_index"]
    arrays["value_targets"][idx] = pos["state_value"]
    arrays["legal_masks"][idx] = legal


# ============================================================================
# Memory-mapped Dataset
# ============================================================================

class MemmappedPositionDataset(Dataset[dict[str, Any]]):
    """Memory-mapped dataset — near-zero RAM usage.

    Loads preprocessed numpy arrays via np.memmap. The OS handles
    paging: only the batch currently being accessed lives in RAM.
    """

    def __init__(self, split_dir: str | Path) -> None:
        split_dir = Path(split_dir)
        meta_path = split_dir.parent / "metadata.json"

        with meta_path.open() as f:
            self.meta = json.load(f)

        split_name = split_dir.name  # "train" or "val"
        count = self.meta[f"{split_name}_positions"]
        max_seq_len = self.meta["max_seq_len"]

        self.input_ids = np.memmap(
            split_dir / "input_ids.dat",
            dtype=np.int16, mode="r", shape=(count, max_seq_len),
        )
        self.attention_mask = np.memmap(
            split_dir / "attention_mask.dat",
            dtype=np.uint8, mode="r", shape=(count, max_seq_len),
        )
        self.action_targets = np.memmap(
            split_dir / "action_targets.dat",
            dtype=np.int16, mode="r", shape=(count,),
        )
        self.value_targets = np.memmap(
            split_dir / "value_targets.dat",
            dtype=np.float32, mode="r", shape=(count,),
        )
        self.legal_masks = np.memmap(
            split_dir / "legal_masks.dat",
            dtype=np.uint8, mode="r", shape=(count, ACTION_SPACE_SIZE),
        )

    def __len__(self) -> int:
        return len(self.action_targets)

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        return {
            "input_ids": torch.from_numpy(
                self.input_ids[idx].astype(np.int64)
            ),
            "attention_mask": torch.from_numpy(
                self.attention_mask[idx].astype(np.int64)
            ),
            "action_target": torch.tensor(
                int(self.action_targets[idx]), dtype=torch.long
            ),
            "value_target": torch.tensor(
                float(self.value_targets[idx]), dtype=torch.float32
            ),
            "legal_mask": torch.from_numpy(
                self.legal_masks[idx].astype(np.float32)
            ),
        }

    @property
    def vocab_size(self) -> int:
        return int(self.meta["vocab_size"])

    @property
    def max_seq_len(self) -> int:
        return int(self.meta["max_seq_len"])

    def stats(self) -> dict[str, Any]:
        """Quick stats from a sample of the dataset."""
        sample_size = min(1000, len(self))
        indices = np.random.default_rng(42).choice(len(self), sample_size, replace=False)
        seq_lens = np.array([
            self.attention_mask[int(i)].sum() for i in indices
        ])
        return {
            "num_positions": len(self),
            "vocab_size": self.vocab_size,
            "max_seq_len": self.max_seq_len,
            "sampled_seq_len_mean": round(float(seq_lens.mean()), 1),
            "sampled_seq_len_median": round(float(np.median(seq_lens)), 0),
            "sampled_seq_len_max": int(seq_lens.max()),
            "sampled_seq_len_p95": round(float(np.percentile(seq_lens, 95)), 0),
        }
