# bclm/tokenizer.py
"""
Tokenizer interface for BCLM models.

Currently backed by TokenMonster. The tokenizer spec (e.g.
``"tokenmonster:english-32000-consistent-v1"``) is read from each
model's ``config.json`` and passed in at construction time — there is
no global configuration.
"""
from __future__ import annotations

import os
import sys
from typing import Dict, Iterator, List, Optional, Set, Tuple

from .config import (
    EOT_TOKEN_ID,
    PAD_TOKEN_ID,
    TURN_TOKEN_ID,
    USER_TOKEN_ID,
    ASSISTANT_TOKEN_ID,
    SYSTEM_TOKEN_ID,
)


# =============================================================================
# Special-token registries
# =============================================================================

SPECIAL_TOKEN_IDS: Set[int] = {
    EOT_TOKEN_ID,
    PAD_TOKEN_ID,
    TURN_TOKEN_ID,
    USER_TOKEN_ID,
    ASSISTANT_TOKEN_ID,
    SYSTEM_TOKEN_ID,
}

SPECIAL_TOKEN_NAMES: Dict[int, str] = {
    EOT_TOKEN_ID:       "<eot>",
    PAD_TOKEN_ID:       "<pad>",
    TURN_TOKEN_ID:      "<turn>",
    USER_TOKEN_ID:      "<user>",
    ASSISTANT_TOKEN_ID: "<assistant>",
    SYSTEM_TOKEN_ID:    "<system>",
}

ROLE_TO_TOKEN: Dict[str, int] = {
    "system":    SYSTEM_TOKEN_ID,
    "user":      USER_TOKEN_ID,
    "assistant": ASSISTANT_TOKEN_ID,
}

TOKEN_TO_ROLE: Dict[int, str] = {v: k for k, v in ROLE_TO_TOKEN.items()}


# =============================================================================
# Spec parsing
# =============================================================================

DEFAULT_TM_VOCAB: str = "english-32000-consistent-v1"


def _parse_tokenizer_spec(spec: str) -> Tuple[str, str]:
    """
    Parse a tokenizer spec string into ``(backend, locator)``.

    Supported forms::

        "tokenmonster:english-32000-consistent-v1"
        "tm:english-32000-consistent-v1"
        "english-32000-consistent-v1"          # assumed TokenMonster
        "/path/to/vocab.vocab"                 # file path → TokenMonster
        "tokenmonster"  /  "tm"  /  "none"     # use default vocab

    Returns:
        ``(backend, locator)`` where *backend* is always lowercased.
    """
    s = (spec or "").strip()
    if not s or s.lower() in {"none", "null", ""}:
        return ("tokenmonster", DEFAULT_TM_VOCAB)

    if ":" in s:
        backend, locator = s.split(":", 1)
        backend = backend.strip().lower()
        locator = locator.strip()
        if backend in {"tokenmonster", "tm"}:
            return ("tokenmonster", locator or DEFAULT_TM_VOCAB)
        return (backend, locator)

    if s.lower() in {"tokenmonster", "tm"}:
        return ("tokenmonster", DEFAULT_TM_VOCAB)

    return ("tokenmonster", s)


# =============================================================================
# BCLMTokenizer
# =============================================================================

class BCLMTokenizer:
    """
    Central tokenizer for BCLM models, backed by TokenMonster.

    Parameters
    ----------
    spec : str
        Tokenizer specification from the model config (the
        ``"tokenizer"`` field). Examples:
        ``"tokenmonster:english-32000-consistent-v1"``,
        ``"english-32000-consistent-v1"``,
        ``"/path/to/custom.vocab"``.
    vocab_size : int
        Total model vocabulary size (including special tokens and any
        reserved padding). This is the embedding / logit dimension,
        **not** the base tokenizer vocabulary size.
    """

    def __init__(self, spec: str, vocab_size: int) -> None:
        self._spec = spec
        self._vocab_size = int(vocab_size)

        backend, locator = _parse_tokenizer_spec(spec)

        if backend != "tokenmonster":
            raise ValueError(
                f"Unsupported tokenizer backend: {backend!r} (spec={spec!r}). "
                "Only 'tokenmonster' is currently supported."
            )

        try:
            import tokenmonster  # type: ignore
        except ImportError as exc:
            raise ImportError(
                "TokenMonster is required but not installed.\n"
                "Install with:  pip install tokenmonster"
            ) from exc

        tm_dir = os.environ.get("BCLM_TOKENMONSTER_DIR", "").strip()
        if tm_dir and hasattr(tokenmonster, "set_local_directory"):
            tokenmonster.set_local_directory(tm_dir)

        self._vocab = tokenmonster.load(locator)
        self._base_vocab_size = int(len(self._vocab))

        # Sanity: base vocab must not overlap special token IDs
        min_special = min(SPECIAL_TOKEN_IDS)
        if self._base_vocab_size > min_special:
            raise ValueError(
                f"TokenMonster vocab size {self._base_vocab_size} overlaps "
                f"reserved special-token IDs (min={min_special}). "
                f"Use a base vocab ≤ {min_special}."
            )

        max_special = max(SPECIAL_TOKEN_IDS)
        if self._vocab_size <= max_special:
            raise ValueError(
                f"Model vocab_size={self._vocab_size} must be > "
                f"max special token id={max_special}."
            )

        if self._vocab_size < self._base_vocab_size:
            raise ValueError(
                f"Model vocab_size={self._vocab_size} < base tokenizer "
                f"vocab size={self._base_vocab_size}."
            )

    # ── Properties ──────────────────────────────────────────

    @property
    def vocab_size(self) -> int:
        """Total vocabulary size including reserved special tokens."""
        return self._vocab_size

    @property
    def base_vocab_size(self) -> int:
        """Base tokenizer vocabulary size (regular tokens only)."""
        return self._base_vocab_size

    @property
    def eot_token_id(self) -> int:
        return EOT_TOKEN_ID

    @property
    def pad_token_id(self) -> int:
        return PAD_TOKEN_ID

    @property
    def user_token_id(self) -> int:
        return USER_TOKEN_ID

    @property
    def assistant_token_id(self) -> int:
        return ASSISTANT_TOKEN_ID

    @property
    def system_token_id(self) -> int:
        return SYSTEM_TOKEN_ID

    # ── Core encode / decode ────────────────────────────────

    def encode(self, text: str) -> List[int]:
        """
        Encode raw text to regular token IDs.

        Does **not** prepend or append special / role tokens — that
        is handled by ``encode_chat``.
        """
        if not text:
            return []
        toks = self._vocab.tokenize(text)
        if hasattr(toks, "tolist"):
            toks = toks.tolist()
        if isinstance(toks, int):
            return [int(toks)]
        return [int(t) for t in toks]

    def decode(self, ids: List[int]) -> str:
        """
        Decode a list of token IDs to text.

        Regular IDs are decoded through TokenMonster; out-of-range IDs
        are rendered as ``<tok:ID>``.
        """
        if not ids:
            return ""

        parts: List[str] = []
        buf: List[int] = []

        def _flush() -> None:
            if buf:
                parts.append(self._vocab.decode(buf))
                buf.clear()

        for tid in ids:
            t = int(tid)
            if 0 <= t < self._base_vocab_size:
                buf.append(t)
            else:
                _flush()
                if t in SPECIAL_TOKEN_NAMES:
                    parts.append(SPECIAL_TOKEN_NAMES[t])
                else:
                    parts.append(f"<tok:{t}>")

        _flush()
        return "".join(parts)

    # ── Chat encoding ───────────────────────────────────────

    def encode_chat(
        self,
        messages: List[Dict[str, str]],
        *,
        add_generation_prompt: bool = True,
    ) -> List[int]:
        """
        Encode a multi-turn conversation into a flat token-ID list.

        Each message is converted to::

            <role_token>  content_ids...

        Completed assistant messages receive an ``<eot>`` suffix.
        When *add_generation_prompt* is True, an ``<assistant>`` token
        is appended so the model begins generating a response.
        """
        tokens: List[int] = []

        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")
            role_token = ROLE_TO_TOKEN.get(role)
            if role_token is None:
                continue

            tokens.append(role_token)
            if content:
                tokens.extend(self.encode(content))

            if role == "assistant":
                tokens.append(EOT_TOKEN_ID)

        if add_generation_prompt:
            if not tokens or tokens[-1] != ASSISTANT_TOKEN_ID:
                tokens.append(ASSISTANT_TOKEN_ID)

        return tokens

    # ── Decoding helpers ────────────────────────────────────

    def decode_response(self, ids: List[int]) -> str:
        """Decode token IDs to text, stripping all special tokens."""
        filtered = [i for i in ids if i not in SPECIAL_TOKEN_IDS]
        if not filtered:
            return ""
        return self.decode(filtered)

    def decode_with_labels(self, ids: List[int]) -> str:
        """
        Decode token IDs, rendering special tokens as readable labels
        (e.g. ``<eot>``, ``<user>``). Useful for debugging.
        """
        parts: List[str] = []
        buf: List[int] = []

        def _flush() -> None:
            if buf:
                parts.append(self.decode(list(buf)))
                buf.clear()

        for tid in ids:
            if tid in SPECIAL_TOKEN_NAMES:
                _flush()
                parts.append(SPECIAL_TOKEN_NAMES[tid])
            else:
                buf.append(tid)
        _flush()
        return "".join(parts)

    # ── Streaming decode ────────────────────────────────────

    def stream_decode(self, token_iter: Iterator[int]) -> Iterator[str]:
        """
        Wrap a token-ID iterator and yield incremental text chunks.

        Handles the common subword edge-case where a single token does
        not map to a clean string boundary by accumulating and
        re-decoding. Special tokens are silently skipped.
        """
        buf: List[int] = []
        prev_text = ""

        for tid in token_iter:
            if tid in SPECIAL_TOKEN_IDS:
                continue
            buf.append(tid)
            try:
                full = self.decode(list(buf))
            except Exception:
                continue
            delta = full[len(prev_text):]
            if delta:
                prev_text = full
                yield delta

        if buf:
            try:
                final = self.decode(buf)
                delta = final[len(prev_text):]
                if delta:
                    yield delta
            except Exception:
                pass

    # ── Validation ──────────────────────────────────────────

    def validate(self) -> None:
        """Run quick sanity checks on encode/decode round-tripping."""
        assert self.vocab_size > 0

        enc = self.encode("hello world")
        assert isinstance(enc, list) and len(enc) > 0
        assert all(isinstance(i, int) for i in enc)

        dec = self.decode(enc)
        assert isinstance(dec, str) and len(dec) > 0

        chat_ids = self.encode_chat([{"role": "user", "content": "hi"}])
        assert chat_ids[0] == USER_TOKEN_ID
        assert chat_ids[-1] == ASSISTANT_TOKEN_ID

    def __repr__(self) -> str:
        return (
            f"BCLMTokenizer(spec={self._spec!r}, "
            f"base_vocab={self._base_vocab_size}, "
            f"model_vocab={self._vocab_size})"
        )