# file: dataset.py
import os
import torch
from torch.utils.data import Dataset
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace


def build_or_load_tokenizer(corpus_path, tokenizer_save_path, vocab_size=1000):
    """Train a BPE tokenizer on the corpus, or load the cached one if present.

    NOTE: the cache is keyed only by path — if you change vocab_size, delete the
    saved json first or you'll silently keep the old vocab.
    """
    if os.path.exists(tokenizer_save_path):
        return Tokenizer.from_file(tokenizer_save_path)

    tokenizer = Tokenizer(BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = Whitespace()  # split on spaces first so merges stay within words
    # "\n" is reserved as a token because each log line is a record boundary
    trainer = BpeTrainer(vocab_size=vocab_size, special_tokens=["<pad>", "<unk>", "\n"])
    tokenizer.train([corpus_path], trainer)
    tokenizer.save(tokenizer_save_path)
    return tokenizer


class AutoregressiveLogDataset(Dataset):
    """Sliding windows over the token stream for next-token training."""

    def __init__(self, corpus_path, tokenizer, max_seq_len=64, stride=None):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        # stride = how far the window advances each step. Default half-window is a
        # balance: stride=1 gives ~187k near-duplicate windows (slow), stride=max_seq_len
        # gives no overlap (less coverage).
        self.stride = stride if stride is not None else max_seq_len // 2

        # Small corpus, so just read + encode it once. Would stream for a big one.
        with open(corpus_path, "r") as f:
            raw_text = f.read()
        self.encoded_tokens = tokenizer.encode(raw_text).ids

        # Precompute window starts; each needs max_seq_len + 1 tokens (the +1 is the last target)
        last_start = len(self.encoded_tokens) - self.max_seq_len - 1
        self.start_positions = list(range(0, max(0, last_start) + 1, self.stride))

    def __len__(self):
        return len(self.start_positions)

    def __getitem__(self, idx):
        start = self.start_positions[idx]
        chunk = self.encoded_tokens[start : start + self.max_seq_len + 1]
        # y is x shifted by one — at every position the model predicts the next token
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y