# Copyright 2026 DoL contributors
# Licensed under the Apache License, Version 2.0.
"""Hugging Face adapter for the versioned DoL dual-track tokenizer.

The adapter deliberately contains no alternative tokenization algorithm.  It
loads the exact MorphBPE/general-BPE assets from ``bundle_v3b`` and exposes the
small interface required by the official Transformers causal-LM example.
"""

from __future__ import annotations

import hashlib
import json
import os
import shutil
from pathlib import Path
from transformers import PreTrainedTokenizer

from Tokenizer.unified.bundle import TokenizerBundle, TokenizerBundleConfig


MODEL_VOCAB_SIZE = 65_536
EXPECTED_SPECIAL_IDS = {
    "<pad>": 0,
    "<unk>": 1,
    "<bos>": 2,
    "<eos>": 3,
    "<img>": 4,
    "<image>": 5,
    "<image_start>": 6,
    "<image_patch>": 7,
    "<image_end>": 8,
    "<video>": 9,
    "<video_start>": 10,
    "<video_patch>": 11,
    "<video_end>": 12,
    "<bbox>": 13,
    "<ocr>": 14,
    "<ocr_start>": 15,
    "<ocr_end>": 16,
    "▁": 17,
    "◈": 18,
    "<doc>": 19,
    "<table>": 20,
    "<layout>": 21,
    "<audio>": 22,
    "<audio_start>": 23,
    "<audio_patch>": 24,
    "<audio_end>": 25,
}


def _sha256(path: str | os.PathLike[str]) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


class DoLTokenizer(PreTrainedTokenizer):
    """Slow HF-compatible facade over the canonical DoL tokenizer runtime.

    Raw language-model text is literal by default: strings that happen to look
    like ``<image>`` are not silently converted into control tokens.  Callers
    that intentionally serialize multimodal control tokens can pass
    ``split_special_tokens=False`` explicitly.
    """

    vocab_files_names = {
        "vocab_file": "vocab.json",
        "morphbpe_file": "morphbpe.json",
        "general_file": "general.json",
        "dol_config_file": "dol_tokenizer_config.json",
        "dol_manifest_file": "dol_tokenizer_manifest.json",
    }
    model_input_names = ["input_ids", "attention_mask"]
    slow_tokenizer_class = None

    def __init__(
        self,
        vocab_file: str,
        morphbpe_file: str,
        general_file: str,
        dol_config_file: str,
        dol_manifest_file: str,
        expected_manifest_sha256: str | None = None,
        model_vocab_size: int = MODEL_VOCAB_SIZE,
        **kwargs,
    ) -> None:
        self.vocab_file = str(vocab_file)
        self.morphbpe_file = str(morphbpe_file)
        self.general_file = str(general_file)
        self.dol_config_file = str(dol_config_file)
        self.dol_manifest_file = str(dol_manifest_file)
        self.model_vocab_size = int(model_vocab_size)
        self.expected_manifest_sha256 = expected_manifest_sha256

        self._verify_assets()
        with open(self.dol_config_file, "r", encoding="utf-8") as handle:
            raw_config = json.load(handle)
        with open(self.vocab_file, "r", encoding="utf-8") as handle:
            vocab = {str(token): int(index) for token, index in json.load(handle).items()}

        config = TokenizerBundleConfig(**raw_config)
        self.bundle = TokenizerBundle._build(
            config,
            morphbpe_path=self.morphbpe_file,
            general_path=self.general_file,
            vocab=vocab,
        )
        self._native = self.bundle.tokenizer
        self._verify_contract()

        kwargs.setdefault("model_max_length", 131_072)
        kwargs.setdefault("padding_side", "right")
        kwargs.setdefault("truncation_side", "right")
        kwargs.setdefault("clean_up_tokenization_spaces", False)
        kwargs.setdefault("split_special_tokens", True)
        kwargs.setdefault("special_tokens_pattern", "bos_eos")
        kwargs.setdefault("token_type_ids_pattern", "all_zeros")
        kwargs.update(
            {
                "pad_token": "<pad>",
                "unk_token": "<unk>",
                "bos_token": "<bos>",
                "eos_token": "<eos>",
                "additional_special_tokens": [
                    token
                    for token, index in sorted(EXPECTED_SPECIAL_IDS.items(), key=lambda item: item[1])
                    if token not in {"<pad>", "<unk>", "<bos>", "<eos>", "▁", "◈"}
                ],
                "vocab_file": self.vocab_file,
                "morphbpe_file": self.morphbpe_file,
                "general_file": self.general_file,
                "dol_config_file": self.dol_config_file,
                "dol_manifest_file": self.dol_manifest_file,
                "expected_manifest_sha256": self.expected_manifest_sha256,
                "model_vocab_size": self.model_vocab_size,
            }
        )
        super().__init__(**kwargs)

    @property
    def vocab_size(self) -> int:
        """Number of assigned IDs; the model keeps a padded 65,536-row table."""

        return self._native.vocab_size

    def __len__(self) -> int:
        """Return max assigned ID + 1, not the number of populated sparse rows."""

        return self.vocab_size

    def get_vocab(self) -> dict[str, int]:
        vocab = dict(self._native.vocab)
        vocab.update(getattr(self, "_added_tokens_encoder", {}))
        return vocab

    def _add_tokens(self, new_tokens, special_tokens: bool = False) -> int:
        """Register existing specials but forbid unversioned vocabulary growth.

        The DoL vocabulary has intentional holes.  The generic HF implementation
        allocates new IDs from ``len(get_vocab())`` and would therefore collide
        with an already assigned sparse ID.  A new token requires a new bundle
        version instead of an in-place ``add_tokens`` call.
        """

        unknown = [str(token) for token in (new_tokens or []) if str(token) not in self._native.vocab]
        if unknown:
            raise ValueError(
                "DoL tokenizer vocabulary is immutable; create a versioned bundle instead of adding tokens: "
                + ", ".join(repr(token) for token in unknown[:8])
            )
        return super()._add_tokens(new_tokens, special_tokens=special_tokens)

    def _tokenize(self, text: str, **kwargs) -> list[str]:
        ids = self._native.encode_plain_text(text)
        return [self._native.id_to_token[index] for index in ids]

    def _convert_token_to_id(self, token: str) -> int:
        return self._native.vocab.get(token, self._native.unk_id)

    def _convert_id_to_token(self, index: int) -> str:
        token = self._native.id_to_token.get(int(index))
        if token is not None:
            return token
        if 0 <= int(index) < self.model_vocab_size:
            return f"<unused_{int(index)}>"
        return "<unk>"

    def convert_tokens_to_string(self, tokens: list[str]) -> str:
        parts: list[str] = []
        assigned: list[int] = []

        def flush() -> None:
            if assigned:
                parts.append(self._native.decode(assigned))
                assigned.clear()

        for token in tokens:
            index = self._native.vocab.get(str(token))
            if index is None:
                flush()
                parts.append(str(token))
            else:
                assigned.append(index)
        flush()
        return "".join(parts)

    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str, ...]:
        if not os.path.isdir(save_directory):
            raise ValueError(f"save_directory is not a directory: {save_directory}")
        prefix = f"{filename_prefix}-" if filename_prefix else ""
        sources = {
            "vocab_file": self.vocab_file,
            "morphbpe_file": self.morphbpe_file,
            "general_file": self.general_file,
            "dol_config_file": self.dol_config_file,
            "dol_manifest_file": self.dol_manifest_file,
        }
        outputs: list[str] = []
        for key, source in sources.items():
            destination = os.path.join(save_directory, prefix + self.vocab_files_names[key])
            if os.path.abspath(source) != os.path.abspath(destination):
                shutil.copyfile(source, destination)
            outputs.append(destination)
        return tuple(outputs)

    def _verify_assets(self) -> None:
        paths = {
            "vocab.json": self.vocab_file,
            "morphbpe.json": self.morphbpe_file,
            "general.json": self.general_file,
            "config.json": self.dol_config_file,
        }
        for name, path in paths.items():
            if not Path(path).is_file():
                raise FileNotFoundError(f"DoL tokenizer asset is missing: {name} ({path})")
        if not Path(self.dol_manifest_file).is_file():
            raise FileNotFoundError(f"DoL tokenizer manifest is missing: {self.dol_manifest_file}")
        if self.expected_manifest_sha256:
            actual_manifest_hash = _sha256(self.dol_manifest_file)
            if actual_manifest_hash != self.expected_manifest_sha256:
                raise ValueError(
                    "DoL tokenizer manifest hash mismatch: "
                    f"expected {self.expected_manifest_sha256}, got {actual_manifest_hash}"
                )
        with open(self.dol_manifest_file, "r", encoding="utf-8") as handle:
            manifest = json.load(handle)
        for name, expected in (manifest.get("files") or {}).items():
            path = paths.get(name)
            if path is None:
                raise ValueError(f"DoL tokenizer manifest contains an unsupported asset: {name}")
            actual = _sha256(path)
            if actual != expected:
                raise ValueError(f"DoL tokenizer asset hash mismatch for {name}: expected {expected}, got {actual}")

    def _verify_contract(self) -> None:
        vocab = self._native.vocab
        values = list(vocab.values())
        if len(values) != len(set(values)):
            raise ValueError("DoL tokenizer vocabulary contains duplicate IDs")
        if not values or min(values) < 0 or max(values) >= self.model_vocab_size:
            raise ValueError(
                f"DoL tokenizer IDs must fit model vocab [0, {self.model_vocab_size}); "
                f"observed min={min(values, default=None)}, max={max(values, default=None)}"
            )
        for token, expected in EXPECTED_SPECIAL_IDS.items():
            actual = vocab.get(token)
            if actual != expected:
                raise ValueError(f"special token {token!r}: expected ID {expected}, got {actual}")
        issues = self.bundle.validate()
        if issues:
            raise ValueError("DoL tokenizer bundle validation failed: " + "; ".join(issues))


__all__ = ["DoLTokenizer", "EXPECTED_SPECIAL_IDS", "MODEL_VOCAB_SIZE"]