from typing import Callable

import torch
from datasets import load_dataset

from .Dataset import Dataset


class EnWik9DataSet(Dataset):
    """
    Hugging Face: https://huggingface.co/datasets/haukur/enwik9
    """
    def __init__(self, root: str | None = None, transform: Callable | None = None):
        super().__init__('enwik9', root, transform)

        # HuggingFace dataset: string text
        data = load_dataset("haukur/enwik9", cache_dir=self.root, split="train")

        # Extract raw text
        text = data["text"]

        # Convert text (Python string) → bytes → tensor of ints 0–255
        # UTF-8 but non-ASCII bytes may exceed 255, so enforce modulo or ignore errors
        byte_data = "".join(text).encode("utf-8", errors="replace")
        self.data = torch.tensor(list(byte_data), dtype=torch.long)

        # Model uses fixed 128-length context
        self.context_length = 128

    def __len__(self):
        # number of sliding windows
        return len(self.data) - self.context_length

    def __getitem__(self, idx):
        # context window
        x = self.data[idx: idx + self.context_length]

        # next byte target
        y = self.data[idx + self.context_length]

        if self.transform:
            x = self.transform(x)

        return x, y
