"""
BoccaccioAI - PyTorch Datasets for Pre-Training and Instruction Fine-Tuning

Provides Dataset classes for loading pre-tokenized binary data (pre-training)
and JSONL instruction data (fine-tuning).

De Lauretis tech
"""

from __future__ import annotations

import json
from pathlib import Path

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


class PreTokenizedDataset(Dataset):
    """Dataset for loading pre-tokenized binary data during pre-training.

    Reads a flat uint16 numpy memmap file produced by ``tokenize_corpus.py``
    and serves fixed-length token sequences.
    """

    def __init__(self, data_path: str, sequence_length: int = 2048) -> None:
        self.data_path = Path(data_path)
        self.sequence_length = sequence_length

        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

        self.data = np.memmap(self.data_path, dtype=np.uint16, mode="r")
        self.num_sequences = len(self.data) // self.sequence_length

        if self.num_sequences == 0:
            raise ValueError(
                f"Data file {self.data_path} contains {len(self.data)} tokens, "
                f"which is fewer than one full sequence of length {self.sequence_length}."
            )

    def __len__(self) -> int:
        return self.num_sequences

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        if idx < 0 or idx >= self.num_sequences:
            raise IndexError(
                f"Index {idx} out of range for dataset with {self.num_sequences} sequences"
            )

        start = idx * self.sequence_length
        end = start + self.sequence_length
        tokens = np.array(self.data[start:end], dtype=np.int64)
        tokens = torch.from_numpy(tokens)

        return {"input_ids": tokens, "labels": tokens}


class InstructionDataset(Dataset):
    """Dataset for instruction fine-tuning with multi-turn conversations.

    Reads JSONL files with the standard format::

        {"conversations": [
            {"role": "system", "content": "..."},     # optional
            {"role": "user", "content": "..."},
            {"role": "assistant", "content": "..."},
            ...
        ]}

    Each conversation is rendered with explicit role markers::

        ### System:
        {system_content}
        ### Utente:
        {user_content}
        ### Assistente:
        {assistant_content}</s>
        ### Utente:
        {user_content_2}
        ### Assistente:
        {assistant_content_2}</s>

    Loss is computed only on assistant tokens (plus the trailing </s>).
    All user, system, and marker tokens are masked with -100.

    Examples longer than ``sequence_length`` after tokenization are skipped
    during ``__init__`` to avoid truncating responses (which would prevent
    the model from learning the EOS token).
    """

    SYSTEM_MARKER = "### System:"
    USER_MARKER = "### Utente:"
    ASSISTANT_MARKER = "### Assistente:"

    def __init__(
        self,
        data_path: str,
        tokenizer_path: str,
        sequence_length: int = 2048,
    ) -> None:
        self.data_path = Path(data_path)
        self.sequence_length = sequence_length

        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

        # Load tokenizer.
        tokenizer_file = Path(tokenizer_path)
        if not tokenizer_file.exists():
            raise FileNotFoundError(f"Tokenizer not found: {tokenizer_file}")
        self.tokenizer = Tokenizer.from_file(str(tokenizer_file))

        # Resolve special token IDs.
        self.pad_token_id = self.tokenizer.token_to_id("<pad>")
        if self.pad_token_id is None:
            self.pad_token_id = 0
        self.eos_token_id = self.tokenizer.token_to_id("</s>")
        if self.eos_token_id is None:
            raise ValueError("EOS token '</s>' not found in tokenizer vocabulary")

        # Load raw examples (JSONL). Tokenization happens on-the-fly in __getitem__
        # to avoid storing 268K pre-tokenized tensors in memory (~13 GB).
        self.raw_examples: list[dict] = []
        with open(self.data_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                self.raw_examples.append(json.loads(line))

        if not self.raw_examples:
            raise ValueError(f"No examples found in {self.data_path}")

        # Pre-compute token lengths to filter out examples that exceed sequence_length.
        # This is fast (just encoding, no tensor creation) and avoids wasting time
        # on examples that would be skipped anyway.
        valid_indices: list[int] = []
        skipped = 0
        for i, ex in enumerate(self.raw_examples):
            token_len = self._estimate_token_length(ex)
            if token_len <= self.sequence_length and token_len >= 10:
                valid_indices.append(i)
            else:
                skipped += 1

        if not valid_indices:
            raise ValueError(
                f"All examples were skipped (too long for sequence_length={sequence_length}). "
                f"Consider increasing sequence_length."
            )

        if skipped > 0:
            print(
                f"  InstructionDataset: skipped {skipped} examples "
                f"(exceeded sequence_length={sequence_length})"
            )

        self.valid_indices = valid_indices

    def _render_turn(self, role: str, content: str) -> str:
        """Render a single conversation turn with its role marker."""
        if role == "system":
            return f"{self.SYSTEM_MARKER}\n{content}\n"
        elif role == "user":
            return f"{self.USER_MARKER}\n{content}\n"
        elif role == "assistant":
            return f"{self.ASSISTANT_MARKER}\n{content}</s>\n"
        else:
            raise ValueError(f"Unknown role: {role}")

    def _estimate_token_length(self, ex: dict) -> int:
        """Exact token count for filtering.

        Mirrors the per-turn tokenization of _process_example exactly: BPE is
        not additive across turn boundaries, so encoding the concatenated text
        in one pass can undercount versus per-turn encoding.
        """
        conversations = ex.get("conversations", [])
        if not conversations or len(conversations) < 2:
            return 0
        if conversations[-1]["role"] != "assistant":
            return 0
        total = 0
        for turn in conversations:
            role = turn.get("role", "")
            content = (turn.get("content") or "").strip()
            if not content:
                continue
            if role in ("system", "user"):
                text = self._render_turn(role, content)
                total += len(self.tokenizer.encode(text).ids)
            elif role == "assistant":
                marker_text = f"{self.ASSISTANT_MARKER}\n"
                total += len(self.tokenizer.encode(marker_text).ids)
                content_text = f"{content}</s>\n"
                total += len(self.tokenizer.encode(content_text).ids)
        return total

    def _process_example(self, ex: dict) -> dict[str, torch.Tensor] | None:
        """Tokenize a conversation example and build labels with masking.

        Returns None if the example is too long for sequence_length.
        """
        conversations = ex.get("conversations", [])
        if not conversations or len(conversations) < 2:
            return None

        # Validate: must end with assistant turn.
        if conversations[-1]["role"] != "assistant":
            return None

        input_ids: list[int] = []
        labels: list[int] = []  # -100 for masked, token_id for assistant content + </s>

        for turn in conversations:
            role = turn.get("role", "")
            content = (turn.get("content") or "").strip()
            if not content:
                continue

            if role == "system":
                # Render and tokenize the marker + content. All masked.
                text = self._render_turn(role, content)
                ids = self.tokenizer.encode(text).ids
                input_ids.extend(ids)
                labels.extend([-100] * len(ids))

            elif role == "user":
                # Render and tokenize the marker + content. All masked.
                text = self._render_turn(role, content)
                ids = self.tokenizer.encode(text).ids
                input_ids.extend(ids)
                labels.extend([-100] * len(ids))

            elif role == "assistant":
                # Render marker (masked) + content + </s> (not masked).
                marker_text = f"{self.ASSISTANT_MARKER}\n"
                marker_ids = self.tokenizer.encode(marker_text).ids

                content_text = f"{content}</s>\n"
                content_ids = self.tokenizer.encode(content_text).ids

                input_ids.extend(marker_ids)
                labels.extend([-100] * len(marker_ids))

                input_ids.extend(content_ids)
                labels.extend(content_ids)

            else:
                # Unknown role, skip.
                continue

        # Check length: skip if too long (don't truncate - would lose </s>).
        if len(input_ids) > self.sequence_length:
            return None
        if len(input_ids) < 10:
            return None

        actual_length = len(input_ids)
        padding_length = self.sequence_length - actual_length

        # Pad input_ids with pad_token_id.
        input_ids_padded = input_ids + [self.pad_token_id] * padding_length

        # Pad labels with -100 (ignored in loss).
        labels_padded = labels + [-100] * padding_length

        # Attention mask: 1 for real tokens, 0 for padding.
        attention_mask = [1] * actual_length + [0] * padding_length

        return {
            "input_ids": torch.tensor(input_ids_padded, dtype=torch.long),
            "labels": torch.tensor(labels_padded, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
        }

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

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        if idx < 0 or idx >= len(self.valid_indices):
            raise IndexError(
                f"Index {idx} out of range for dataset with {len(self.valid_indices)} examples"
            )
        raw_idx = self.valid_indices[idx]
        ex = self.raw_examples[raw_idx]
        processed = self._process_example(ex)
        # processed should not be None since we pre-filtered, but guard anyway.
        if processed is None:
            # Fallback: return a minimal valid example (shouldn't happen).
            raise RuntimeError(
                f"Example at index {idx} (raw {raw_idx}) failed to process "
                f"despite passing pre-filter. This is a bug."
            )
        return processed


if __name__ == "__main__":
    import sys

    print("BoccaccioAI Dataset classes")
    print(f"  PreTokenizedDataset: for pre-training on binary memmap files")
    print(f"  InstructionDataset:  for instruction fine-tuning on JSONL data")

    if len(sys.argv) > 1 and sys.argv[1] == "--test-pretokenized":
        if len(sys.argv) < 3:
            print("Usage: python dataset.py --test-pretokenized <data_path>")
            sys.exit(1)
        ds = PreTokenizedDataset(sys.argv[2])
        print(f"  Loaded {len(ds)} sequences from {sys.argv[2]}")
        sample = ds[0]
        print(f"  Sample input_ids shape: {sample['input_ids'].shape}")
        print(f"  Sample input_ids dtype: {sample['input_ids'].dtype}")
        print(f"  First 20 token IDs: {sample['input_ids'][:20].tolist()}")

    if len(sys.argv) > 1 and sys.argv[1] == "--test-instruction":
        if len(sys.argv) < 4:
            print("Usage: python dataset.py --test-instruction <data_path> <tokenizer_path>")
            sys.exit(1)
        ds = InstructionDataset(sys.argv[2], sys.argv[3])
        print(f"  Loaded {len(ds)} examples from {sys.argv[2]}")
        sample = ds[0]
        print(f"  Sample input_ids shape: {sample['input_ids'].shape}")
        print(f"  Sample labels shape:    {sample['labels'].shape}")
        print(f"  Sample attn_mask shape: {sample['attention_mask'].shape}")
        num_masked = (sample["labels"] == -100).sum().item()
        num_active = sample["labels"].numel() - num_masked
        print(f"  Masked label positions: {num_masked}")
        print(f"  Active label positions: {num_active} (assistant tokens + </s>)")