"""Pre-encoded sharded datasets generated by the Rust virtual dataset tool."""

from __future__ import annotations

import bisect
import json
from pathlib import Path
from typing import Dict, List, Optional

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


class ShardedEncodedDataset(Dataset):
    """Map-style dataset backed by pre-encoded `.npy` shards.

    The Rust generator writes compact uint16/int16/u8 arrays. This class loads
    one shard at a time and relies on sequential sampling over pre-shuffled
    shards, so Python does no tokenization or BIO permutation during training.
    """

    preserve_order = True

    def __init__(self, dataset_dir: str | Path, manifest_name: str = "manifest.json"):
        self.dataset_dir = Path(dataset_dir)
        self.manifest_path = self.dataset_dir / manifest_name
        self.manifest = json.loads(self.manifest_path.read_text(encoding="utf-8"))
        if self.manifest.get("format") != "anifilebert.virtual_dataset.shards.v1":
            raise ValueError(f"Unsupported virtual dataset manifest: {self.manifest_path}")

        self.max_length = int(self.manifest["max_length"])
        self.shards: List[Dict] = list(self.manifest.get("shards") or [])
        if not self.shards:
            raise ValueError(f"Virtual dataset has no shards: {self.manifest_path}")

        self._starts: List[int] = []
        total = 0
        for shard in self.shards:
            self._starts.append(total)
            total += int(shard["rows"])
        self.total_rows = total

        declared_total = int(self.manifest.get("total_rows", total))
        if declared_total != total:
            raise ValueError(
                f"Virtual dataset row count mismatch: manifest total_rows={declared_total}, "
                f"shard rows={total}"
            )

        self._cache_index: Optional[int] = None
        self._cache: Optional[Dict[str, np.ndarray]] = None

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        if idx < 0:
            idx += self.total_rows
        if idx < 0 or idx >= self.total_rows:
            raise IndexError(idx)

        shard_idx = bisect.bisect_right(self._starts, idx) - 1
        shard_start = self._starts[shard_idx]
        row_idx = idx - shard_start
        cache = self._load_shard(shard_idx)
        return {
            "input_ids": torch.from_numpy(cache["input_ids"][row_idx]),
            "attention_mask": torch.from_numpy(cache["attention_mask"][row_idx]),
            "labels": torch.from_numpy(cache["labels"][row_idx]),
        }

    def _load_shard(self, shard_idx: int) -> Dict[str, np.ndarray]:
        if self._cache_index == shard_idx and self._cache is not None:
            return self._cache

        shard = self.shards[shard_idx]
        cache = {
            "input_ids": np.load(self.dataset_dir / shard["input_ids"], allow_pickle=False),
            "attention_mask": np.load(self.dataset_dir / shard["attention_mask"], allow_pickle=False),
            "labels": np.load(self.dataset_dir / shard["labels"], allow_pickle=False),
        }
        expected_shape = (int(shard["rows"]), self.max_length)
        for key, array in cache.items():
            if array.shape != expected_shape:
                raise ValueError(
                    f"Shard {shard_idx} {key} has shape {array.shape}, expected {expected_shape}"
                )
        self._cache_index = shard_idx
        self._cache = cache
        return cache


class DatasetRangeView(Dataset):
    """A contiguous range view over another dataset."""

    preserve_order = True

    def __init__(self, dataset: Dataset, start: int, end: int):
        if start < 0 or end < start or end > len(dataset):
            raise ValueError(f"Invalid dataset range [{start}, {end}) for length {len(dataset)}")
        self.dataset = dataset
        self.start = start
        self.end = end

    def __len__(self) -> int:
        return self.end - self.start

    def __getitem__(self, idx: int):
        if idx < 0:
            idx += len(self)
        if idx < 0 or idx >= len(self):
            raise IndexError(idx)
        return self.dataset[self.start + idx]