"""Packed token corpora for STRATA language-model pretraining.

Documents are tokenised, separated by an end-of-sequence token, and packed into
a single flat token array on disk (``tokens.bin``) with a JSON sidecar
(``meta.json``). Training reads fixed-length, non-overlapping windows via
:class:`PackedLMDataset` and a memory map, so multi-billion-token corpora never
need to fit in RAM.

Tokenisation is passed in as a callable returning token ids, which keeps this
module independent of any specific tokenizer and lets the byte tokenizer use a
fast vectorised path instead of building span lattices per byte.
"""

from __future__ import annotations

import hashlib
import json
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterable

import numpy as np

TokenizeIds = Callable[[str], np.ndarray]

TOKENS_FILENAME = "tokens.bin"
META_FILENAME = "meta.json"


def choose_token_dtype(vocab_size: int) -> np.dtype:
    """Smallest unsigned integer dtype that can hold every id in the vocab."""

    if vocab_size <= 0:
        raise ValueError("vocab_size must be positive")
    if vocab_size <= np.iinfo(np.uint16).max + 1:
        return np.dtype(np.uint16)
    if vocab_size <= np.iinfo(np.uint32).max + 1:
        return np.dtype(np.uint32)
    raise ValueError(f"vocab_size {vocab_size} too large for uint32 token ids")


@dataclass(slots=True)
class PackedCorpusMeta:
    """Sidecar metadata describing a packed token corpus."""

    num_tokens: int
    dtype: str
    vocab_size: int
    eos_id: int
    tokenizer_name: str
    documents: int
    per_language_tokens: dict[str, int]
    sha256: str
    seed: int | None = None
    source: dict[str, object] = field(default_factory=dict)

    def to_dict(self) -> dict[str, object]:
        return {
            "num_tokens": self.num_tokens,
            "dtype": self.dtype,
            "vocab_size": self.vocab_size,
            "eos_id": self.eos_id,
            "tokenizer_name": self.tokenizer_name,
            "documents": self.documents,
            "per_language_tokens": self.per_language_tokens,
            "sha256": self.sha256,
            "seed": self.seed,
            "source": self.source,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "PackedCorpusMeta":
        return cls(
            num_tokens=int(data["num_tokens"]),
            dtype=str(data["dtype"]),
            vocab_size=int(data["vocab_size"]),
            eos_id=int(data["eos_id"]),
            tokenizer_name=str(data.get("tokenizer_name", "")),
            documents=int(data.get("documents", 0)),
            per_language_tokens=dict(data.get("per_language_tokens", {}) or {}),  # type: ignore[arg-type]
            sha256=str(data.get("sha256", "")),
            seed=data.get("seed"),  # type: ignore[arg-type]
            source=dict(data.get("source", {}) or {}),  # type: ignore[arg-type]
        )

    def write(self, out_dir: Path) -> Path:
        path = out_dir / META_FILENAME
        path.write_text(json.dumps(self.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
        return path

    @classmethod
    def read(cls, out_dir: Path) -> "PackedCorpusMeta":
        return cls.from_dict(json.loads((out_dir / META_FILENAME).read_text(encoding="utf-8")))


def build_packed_corpus(
    *,
    documents: Iterable[tuple[str, str]],
    tokenize_ids: TokenizeIds,
    eos_id: int,
    vocab_size: int,
    tokenizer_name: str,
    out_dir: Path,
    max_tokens: int | None = None,
    per_language_tokens: dict[str, int] | None = None,
    seed: int | None = None,
    source: dict[str, object] | None = None,
    flush_every_tokens: int = 8_000_000,
    logger=None,
) -> PackedCorpusMeta:
    """Tokenise ``(language, text)`` documents and pack them into ``tokens.bin``.

    ``per_language_tokens`` caps how many tokens each language contributes (for
    balanced multilingual mixtures); ``max_tokens`` caps the total. Both are
    honoured mid-document at token granularity.
    """

    out_dir.mkdir(parents=True, exist_ok=True)
    dtype = choose_token_dtype(vocab_size)
    if not (0 <= eos_id < vocab_size):
        raise ValueError(f"eos_id {eos_id} out of range for vocab_size {vocab_size}")

    tokens_path = out_dir / TOKENS_FILENAME
    caps = dict(per_language_tokens or {})
    lang_counts: Counter[str] = Counter()
    hasher = hashlib.sha256()
    total = 0
    docs = 0
    buffer: list[np.ndarray] = []
    buffered = 0

    def flush() -> None:
        nonlocal buffered
        if not buffer:
            return
        chunk = np.concatenate(buffer)
        chunk.tofile(handle)
        hasher.update(chunk.tobytes())
        buffer.clear()
        buffered = 0

    with tokens_path.open("wb") as handle:
        for language, text in documents:
            if max_tokens is not None and total >= max_tokens:
                break
            if caps and all(lang_counts[lang] >= cap for lang, cap in caps.items()):
                break
            if language in caps and lang_counts[language] >= caps[language]:
                continue
            if not text:
                continue
            ids = np.asarray(tokenize_ids(text), dtype=dtype)
            if ids.size == 0:
                continue
            ids = np.append(ids, np.asarray([eos_id], dtype=dtype))

            # Enforce caps precisely at token granularity.
            if language in caps:
                remaining = caps[language] - lang_counts[language]
                if remaining <= 0:
                    continue
                if ids.size > remaining:
                    ids = ids[:remaining]
            if max_tokens is not None and total + ids.size > max_tokens:
                ids = ids[: max_tokens - total]
            if ids.size == 0:
                continue
            if ids.max() >= vocab_size:
                raise ValueError(
                    f"token id {int(ids.max())} >= vocab_size {vocab_size}; "
                    "tokenizer/model vocab mismatch"
                )

            buffer.append(ids)
            buffered += ids.size
            total += ids.size
            lang_counts[language] += int(ids.size)
            docs += 1
            if buffered >= flush_every_tokens:
                flush()
                if logger is not None:
                    logger.info("packed %d tokens (%d docs)", total, docs)
        flush()

    meta = PackedCorpusMeta(
        num_tokens=total,
        dtype=dtype.name,
        vocab_size=vocab_size,
        eos_id=eos_id,
        tokenizer_name=tokenizer_name,
        documents=docs,
        per_language_tokens=dict(lang_counts),
        sha256=hasher.hexdigest(),
        seed=seed,
        source=dict(source or {}),
    )
    meta.write(out_dir)
    if logger is not None:
        logger.info("packed corpus: %d tokens, %d docs -> %s", total, docs, tokens_path)
    return meta


class PackedLMDataset:
    """Fixed-length, non-overlapping windows over a packed token corpus.

    Implements the ``torch.utils.data.Dataset`` protocol (``__len__`` /
    ``__getitem__``) without importing torch at module load; each item is a 1-D
    ``int64`` tensor of length ``seq_len`` suitable for causal-LM training where
    the model shifts inputs and labels internally.
    """

    def __init__(self, corpus_dir: str | Path, *, seq_len: int) -> None:
        if seq_len <= 1:
            raise ValueError("seq_len must be > 1")
        self.corpus_dir = Path(corpus_dir)
        self.meta = PackedCorpusMeta.read(self.corpus_dir)
        self.seq_len = seq_len
        self.dtype = np.dtype(self.meta.dtype)
        tokens_path = self.corpus_dir / TOKENS_FILENAME
        if not tokens_path.exists():
            raise FileNotFoundError(f"packed tokens not found: {tokens_path}")
        self._tokens = np.memmap(tokens_path, dtype=self.dtype, mode="r")
        self.num_windows = (len(self._tokens)) // seq_len
        if self.num_windows == 0:
            raise ValueError(
                f"corpus has {len(self._tokens)} tokens, too few for seq_len {seq_len}"
            )

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

    def __getitem__(self, index: int):
        import torch

        if index < 0:
            index += self.num_windows
        if not 0 <= index < self.num_windows:
            raise IndexError(index)
        start = index * self.seq_len
        window = np.asarray(self._tokens[start : start + self.seq_len], dtype=np.int64)
        return torch.from_numpy(window)

    def total_tokens_used(self) -> int:
        return self.num_windows * self.seq_len