"""Read full sentence token sequences using read-only memory maps."""

import json
import mmap
import struct
from pathlib import Path


class TokenStore:
    def __init__(self, directory, split="train"):
        self.directory = Path(directory)
        self._resources = []
        if split not in {"train", "validation", "test"}:
            raise ValueError("Unknown split.")
        self.manifest = json.loads((self.directory / "manifest.json").read_text(encoding="utf-8"))
        if self.manifest.get("status") != "complete" or self.manifest.get("format") != "vimeml_sentence_tokens_v1":
            raise ValueError("Incomplete or unsupported token store.")
        if (self.manifest.get("token_dtype"), self.manifest.get("offset_dtype"), self.manifest.get("offset_unit")) != ("uint16_le", "uint64_le", "tokens"):
            raise ValueError("Unsupported token or offset format.")
        self.split = split
        stats = self.manifest["splits"][split]
        self.count = stats["sentences"]
        self.total_tokens = stats["stored_tokens"]
        try:
            self._tokens = self._map(f"{split}.tokens.bin", self.total_tokens * 2)
            self._offsets = self._map(f"{split}.offsets.bin", (self.count + 1) * 8)
            if self._offset(0) != 0 or self._offset(self.count) != self.total_tokens:
                raise ValueError("Invalid initial or final sentence offset.")
        except BaseException:
            self.close()
            raise

    def _map(self, name, expected_size):
        path = self.directory / name
        if path.stat().st_size != expected_size:
            raise ValueError(f"Unexpected file size: {name}")
        stream = path.open("rb")
        self._resources.append(stream)
        if expected_size == 0:
            return None
        mapped = mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_READ)
        self._resources.append(mapped)
        return mapped

    def _offset(self, index):
        return struct.unpack_from("<Q", self._offsets, index * 8)[0]

    def __len__(self):
        return self.count

    def __getitem__(self, index):
        if not self._resources:
            raise ValueError("TokenStore is closed.")
        if not isinstance(index, int):
            raise TypeError("Sentence index must be an integer.")
        if index < 0:
            index += self.count
        if not 0 <= index < self.count:
            raise IndexError(index)
        start, end = self._offset(index), self._offset(index + 1)
        if not 0 <= start < end <= self.total_tokens:
            raise ValueError("Invalid sentence offset range.")
        ids = [item[0] for item in struct.iter_unpack("<H", self._tokens[start * 2:end * 2])]
        special = self.manifest["special_ids"]
        if len(ids) < 3 or ids[0] != special["bos"] or ids[-1] != special["eos"]:
            raise ValueError("Invalid sentence boundary tokens.")
        if any(token >= self.manifest["vocab_size"] for token in ids):
            raise ValueError("Token ID outside vocabulary.")
        return ids

    def close(self):
        while self._resources:
            self._resources.pop().close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()