"""Tokenizer wrapper with strict identity tracking.

Engram hashes token ids, so a tokenizer swap silently rehashes every table row
(docs/ENGRAM_NOTES.md section 1.2). The checkpoint therefore stores this
tokenizer's `identity()`, and `checkpoint.load()` refuses to resume when it
differs. This is the "tokenizer frozen once a run starts" rule from the brief,
enforced rather than documented.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

__all__ = ["Tokenizer", "train_tokenizer", "SPECIAL_TOKENS"]

# ids 0..3; pad is 0 so Engram's pad fill is a real, learnable row.
SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]
PAD_ID, BOS_ID, EOS_ID, UNK_ID = 0, 1, 2, 3


class Tokenizer:
    """Thin wrapper over a HuggingFace `tokenizers.Tokenizer`."""

    def __init__(self, path: str | Path):
        from tokenizers import Tokenizer as HFTokenizer

        self.path = Path(path)
        tok_file = self.path / "tokenizer.json" if self.path.is_dir() else self.path
        if not tok_file.exists():
            raise FileNotFoundError(
                f"no tokenizer.json at {tok_file}. Train one with:\n"
                f"  python scripts/train_tokenizer.py --input <corpus> --out {self.path}"
            )
        self._tok = HFTokenizer.from_file(str(tok_file))
        self._blob = tok_file.read_bytes()
        self.pad_id, self.bos_id, self.eos_id, self.unk_id = PAD_ID, BOS_ID, EOS_ID, UNK_ID
        # Verified, not assumed: a tokenizer.json with a different special-token
        # layout would otherwise inject the wrong id at every document boundary,
        # silently. Only tokens the file actually declares are checked, so a
        # vocabulary without e.g. <unk> still loads.
        for token, expected in zip(SPECIAL_TOKENS, (PAD_ID, BOS_ID, EOS_ID, UNK_ID)):
            actual = self._tok.token_to_id(token)
            if actual is not None and actual != expected:
                raise ValueError(
                    f"{tok_file} maps {token!r} to id {actual}, but this "
                    f"codebase hardcodes id {expected} (packed shards and "
                    f"Engram hashes depend on it)"
                )

    # -- HF-ish surface, enough for engram's CompressedTokenBackend ---------- #

    def __len__(self) -> int:
        return self._tok.get_vocab_size()

    @property
    def vocab_size(self) -> int:
        return self._tok.get_vocab_size()

    @property
    def backend_tokenizer(self):
        return self._tok

    def encode(self, text: str, add_special: bool = False) -> list[int]:
        ids = self._tok.encode(text, add_special_tokens=False).ids
        if add_special:
            ids = [self.bos_id, *ids, self.eos_id]
        return ids

    def decode(self, ids, skip_special_tokens: bool = True) -> str:
        return self._tok.decode(list(ids), skip_special_tokens=skip_special_tokens)

    def id_to_token(self, i: int):
        return self._tok.id_to_token(i)

    # -- identity ----------------------------------------------------------- #

    def identity(self) -> dict:
        """Everything needed to prove two runs used the same tokenizer."""
        return {
            "sha256": hashlib.sha256(self._blob).hexdigest(),
            "vocab_size": self.vocab_size,
            "pad_id": self.pad_id,
            "bos_id": self.bos_id,
            "eos_id": self.eos_id,
        }

    def save_into(self, directory: str | Path) -> None:
        """Copy the tokenizer into a checkpoint directory, so the checkpoint is
        self-contained and a resume cannot pick up a drifted file."""
        directory = Path(directory)
        directory.mkdir(parents=True, exist_ok=True)
        (directory / "tokenizer.json").write_bytes(self._blob)
        (directory / "tokenizer_identity.json").write_text(
            json.dumps(self.identity(), indent=2), encoding="utf-8"
        )


def train_tokenizer(
    files: list[str | Path],
    out_dir: str | Path,
    vocab_size: int = 32768,
    min_frequency: int = 2,
) -> Tokenizer:
    """Train a byte-level BPE and write `tokenizer.json`.

    Byte-level so the vocabulary can never produce <unk> on unseen text, which
    keeps Engram's hashing total.
    """
    from tokenizers import Tokenizer as HFTokenizer
    from tokenizers import decoders, models, pre_tokenizers, trainers

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tok = HFTokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=SPECIAL_TOKENS,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    )
    tok.train([str(f) for f in files], trainer)
    tok.save(str(out_dir / "tokenizer.json"))

    wrapper = Tokenizer(out_dir)
    wrapper.save_into(out_dir)
    return wrapper