import torch
from torch.utils.data import Dataset
from model.config import ModelConfig
# We assume tokenizer is passed as an instance or we use generic types
# Avoiding direct import of Tokenizer here if we pass the instance to avoid circular deps if any,
# but importing for type hinting is usually fine.
from model.tokenizer import AdvancedTokenizer

class TextDataset(Dataset):
    def __init__(self, text: str, tokenizer: AdvancedTokenizer, config: ModelConfig, pin_memory: bool = False):
        """
        Args:
            text (str): The full training corpus.
            tokenizer (AdvancedTokenizer): The tokenizer instance.
            config (ModelConfig): Configuration object.
            pin_memory (bool): If True, tensor becomes DMA-locked in system RAM for faster GPU transfer.
        """
        print("Encoding dataset (this may take a moment)...")
        self.data = torch.tensor(tokenizer.encode(text), dtype=torch.long)
        
        if pin_memory and torch.cuda.is_available():
            print("Pinning Dataset to Fixed RAM (Pinned Memory)...")
            self.data = self.data.pin_memory()
            
        self.block_size = config.block_size
        print(f"Dataset loaded. Total tokens: {len(self.data)}")

    def __len__(self):
        # We need block_size + 1 characters for input(x) and target(y)
        # So valid start indices are up to len - block_size
        if len(self.data) <= self.block_size:
            return 0
        return len(self.data) - self.block_size

    def __getitem__(self, idx):
        # Grab a chunk of (block_size + 1) tokens
        chunk = self.data[idx : idx + self.block_size + 1]
        
        # x is inputs, y is targets (shifted by 1)
        x = chunk[:-1]
        y = chunk[1:]
        
        return x, y