# SPDX-License-Identifier: Apache-2.0
"""
Tokenizer utilities for oMLX.

This module provides shared tokenizer configuration and fixes that are used
across multiple modules in the codebase.
"""

import json
import logging
from collections.abc import Callable
from functools import lru_cache, partial
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


def unwrap_tokenizer(tokenizer):
    """Unwrap mlx-lm TokenizerWrapper to a HuggingFace PreTrainedTokenizer.

    xgrammar accepts HuggingFace ``PreTrainedTokenizer`` /
    ``PreTrainedTokenizerFast`` but NOT the raw ``tokenizers.Tokenizer``
    nor the mlx-lm ``TokenizerWrapper``.  This helper peels exactly one
    layer of mlx-lm wrapping while keeping the HuggingFace object intact.
    """
    try:
        from transformers import PreTrainedTokenizerBase

        if isinstance(tokenizer, PreTrainedTokenizerBase):
            return tokenizer
    except ImportError:
        pass
    if hasattr(tokenizer, "_tokenizer"):
        inner = tokenizer._tokenizer
        try:
            from transformers import PreTrainedTokenizerBase

            if isinstance(inner, PreTrainedTokenizerBase):
                return inner
        except ImportError:
            pass
        return inner
    return tokenizer


def resolve_vocab_size(model: Any) -> int | None:
    """Extract vocab_size from a model's config/args, handling nested configs.

    Tries ``model.config.vocab_size``, then ``model.args.vocab_size``,
    then ``text_config.vocab_size`` for VLM composite models (e.g. Qwen3.5).

    Args:
        model: An MLX model object (LLM, VLM, or any object with config/args).

    Returns:
        The vocabulary size, or None if it cannot be determined.
    """
    if model is None:
        return None
    for attr in ("config", "args"):
        config = getattr(model, attr, None)
        if config is None:
            continue
        vs = getattr(config, "vocab_size", None)
        if isinstance(vs, int):
            return vs
        text_cfg = getattr(config, "text_config", None)
        if isinstance(text_cfg, dict):
            vs = text_cfg.get("vocab_size")
        elif text_cfg is not None:
            vs = getattr(text_cfg, "vocab_size", None)
        if isinstance(vs, int):
            return vs
    return None


def is_harmony_model(model_name: str, config: dict[str, Any] | None = None) -> bool:
    """
    Check if the model uses Harmony format.

    Harmony format is used by gpt-oss models with special tokens like
    <|start|>, <|channel|>, <|message|>, <|end|>, <|return|>, <|call|>.

    Detection priority:
    1. model_type == "gpt_oss" in config.json
    2. Fallback: model_name contains "gpt-oss" or "gptoss" (case-insensitive)

    Args:
        model_name: The model name or path.
        config: Optional model config dict (from config.json).

    Returns:
        True if the model uses Harmony format.
    """
    # Primary detection: config.model_type
    if config is not None:
        model_type = config.get("model_type", "")
        if model_type == "gpt_oss":
            logger.debug(f"Harmony model detected via config.model_type: {model_name}")
            return True

    # Fallback detection: model name pattern
    if model_name:
        name_lower = model_name.lower()
        if "gpt-oss" in name_lower or "gptoss" in name_lower:
            logger.debug(f"Harmony model detected via model name pattern: {model_name}")
            return True

    return False


def is_gemma4_model(model_name: str, config: dict[str, Any] | None = None) -> bool:
    """
    Check if the model is a Gemma 4 model.

    Detection priority:
    1. Gemma 4 model_type in config.json
    2. Fallback: model_name contains "gemma-4" or "gemma4" (case-insensitive)
    """
    if config is not None:
        model_type = config.get("model_type", "")
        # diffusion_gemma shares Gemma 4's wire protocol (channel markers,
        # call:name{...} tool calls), so it uses the same parser/extractor.
        if model_type in {"gemma4", "gemma4_unified", "diffusion_gemma"}:
            logger.debug(f"Gemma 4 model detected via config.model_type: {model_name}")
            return True

    if model_name:
        name_lower = model_name.lower()
        if "gemma-4" in name_lower or "gemma4" in name_lower:
            logger.debug(f"Gemma 4 model detected via model name pattern: {model_name}")
            return True

    return False


def is_qwen3_model(model_name: str) -> bool:
    """
    Check if the model is a Qwen3 model.

    Args:
        model_name: The model name or path.

    Returns:
        True if the model is a Qwen3 model.
    """
    model_lower = model_name.lower()
    return "qwen3" in model_lower or "Qwen3" in model_name


def _read_json_file(path: Path) -> dict[str, Any] | None:
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError) as exc:
        logger.debug("Failed to read %s: %s", path, exc)
        return None
    return data if isinstance(data, dict) else None


def _find_tokenizer_json(
    tokenizer: Any,
    model_path: str | Path | None = None,
) -> Path | None:
    candidates: list[str | Path] = []
    if model_path:
        candidates.append(model_path)

    tokenizer_path = getattr(tokenizer, "name_or_path", None)
    if tokenizer_path:
        candidates.append(tokenizer_path)

    for candidate in candidates:
        candidate_path = Path(candidate).expanduser()
        tokenizer_file = candidate_path / "tokenizer.json"
        if tokenizer_file.exists():
            return tokenizer_file

        try:
            from huggingface_hub import try_to_load_from_cache

            cached = try_to_load_from_cache(str(candidate), "tokenizer.json")
        except Exception:
            cached = None

        if cached and isinstance(cached, str):
            cached_path = Path(cached)
            if cached_path.exists():
                return cached_path

    return None


@lru_cache(maxsize=128)
def _detokenizer_factory_from_tokenizer_json(
    tokenizer_file: str,
) -> Callable[[Any], Any] | None:
    tokenizer_content = _read_json_file(Path(tokenizer_file))
    if not tokenizer_content or "decoder" not in tokenizer_content:
        return None

    try:
        from mlx_lm.tokenizer_utils import (
            BPEStreamingDetokenizer,
            SPMStreamingDetokenizer,
            _is_bpe_decoder,
            _is_spm_decoder,
            _is_spm_decoder_no_space,
        )
    except ImportError:
        return None

    decoder = tokenizer_content["decoder"]
    if _is_spm_decoder(decoder):
        return SPMStreamingDetokenizer
    if _is_spm_decoder_no_space(decoder):
        return partial(SPMStreamingDetokenizer, trim_space=False)
    if _is_bpe_decoder(decoder):
        return BPEStreamingDetokenizer
    return None


def _is_unsafe_mlx_vlm_bpe_detokenizer(detokenizer: Any) -> bool:
    detokenizer_type = type(detokenizer)
    return (
        detokenizer_type.__module__ == "mlx_vlm.tokenizer_utils"
        and detokenizer_type.__name__ == "BPEStreamingDetokenizer"
    )


def _create_decoder_aware_detokenizer(
    tokenizer: Any,
    tokenizer_file: Path | None,
) -> Any | None:
    if tokenizer_file is None:
        return None

    factory = _detokenizer_factory_from_tokenizer_json(str(tokenizer_file))
    if factory is None:
        return None

    try:
        return factory(tokenizer)
    except Exception as exc:
        logger.debug(
            "Failed to create decoder-aware detokenizer from %s: %s",
            tokenizer_file,
            exc,
        )
        return None


class _CompatNaiveStreamingDetokenizer:
    """Naive fallback for raw tokenizers that lack mlx-lm's probe APIs."""

    def __init__(self, tokenizer: Any):
        self._tokenizer = tokenizer
        self._tokenizer.decode([0])
        self.reset()

    def reset(self) -> None:
        self.offset = 0
        self.tokens = []
        self._text = ""
        self._current_tokens = []
        self._current_text = ""

    def add_token(self, token: int) -> None:
        self._current_tokens.append(token)
        self.tokens.append(token)

    def finalize(self) -> None:
        self._text += self._tokenizer.decode(self._current_tokens)
        self._current_tokens = []
        self._current_text = ""

    @property
    def text(self) -> str:
        if self._current_tokens:
            self._current_text = self._tokenizer.decode(self._current_tokens)
            if self._current_text.endswith("\ufffd") or (
                bool(getattr(self._tokenizer, "clean_up_tokenization_spaces", False))
                and len(self._current_text) > 0
                and self._current_text[-1] == " "
            ):
                self._current_text = self._current_text[:-1]
        if self._current_text and self._current_text[-1] == "\n":
            self._text += self._current_text
            self._current_tokens.clear()
            self._current_text = ""
        return self._text + self._current_text

    @property
    def last_segment(self) -> str:
        text = self.text
        segment = text[self.offset :]
        self.offset = len(text)
        return segment


def create_streaming_detokenizer(
    tokenizer: Any,
    model_path: str | Path | None = None,
) -> Any | None:
    """Create a fresh streaming detokenizer for one request.

    mlx-lm's TokenizerWrapper exposes the correct per-model detokenizer, but
    raw VLM/DFlash tokenizers may not.  In that case, mirror mlx-lm's
    tokenizer.json decoder detection before falling back to the naive decoder.
    """
    has_existing_attr = True
    try:
        detokenizer = tokenizer.detokenizer
    except AttributeError:
        has_existing_attr = False
        detokenizer = None
    except Exception as exc:
        has_existing_attr = False
        detokenizer = None
        logger.debug("Failed to read tokenizer.detokenizer: %s", exc)

    if detokenizer is not None:
        if _is_unsafe_mlx_vlm_bpe_detokenizer(detokenizer):
            tokenizer_file = _find_tokenizer_json(tokenizer, model_path)
            decoder_aware_detokenizer = _create_decoder_aware_detokenizer(
                tokenizer,
                tokenizer_file,
            )
            if decoder_aware_detokenizer is not None:
                return decoder_aware_detokenizer
            logger.debug(
                "Using existing mlx-vlm BPE detokenizer because no "
                "decoder-aware replacement is available"
            )
        return detokenizer

    tokenizer_file = _find_tokenizer_json(tokenizer, model_path)
    decoder_aware_detokenizer = _create_decoder_aware_detokenizer(
        tokenizer,
        tokenizer_file,
    )
    if decoder_aware_detokenizer is not None:
        return decoder_aware_detokenizer

    if has_existing_attr:
        return None

    try:
        from mlx_lm.tokenizer_utils import NaiveStreamingDetokenizer
    except ImportError:
        return None

    try:
        return NaiveStreamingDetokenizer(tokenizer)
    except Exception as exc:
        logger.debug("Failed to create naive streaming detokenizer: %s", exc)

    try:
        return _CompatNaiveStreamingDetokenizer(tokenizer)
    except Exception as compat_exc:
        logger.debug(
            "Failed to create compatibility naive streaming detokenizer: %s",
            compat_exc,
        )
        return None


def _is_lfm2_text_lm(model_name: str) -> bool:
    """Return True for local LFM2 text causal LM checkpoints."""
    config_path = Path(model_name) / "config.json"
    config = _read_json_file(config_path)
    if config is None:
        return False

    model_type = str(config.get("model_type") or "").lower().replace("-", "_")
    architectures = [
        str(arch) for arch in config.get("architectures", []) if isinstance(arch, str)
    ]
    architectures_lower = [arch.lower() for arch in architectures]

    if model_type in {"lfm_audio", "lfm2_audio"}:
        return False
    if any(key in config for key in ("audio_config", "tts_config", "stt_config")):
        return False
    if any("audio" in arch for arch in architectures_lower):
        return False
    if not any("forcausallm" in arch for arch in architectures_lower):
        return False

    return model_type.startswith("lfm2") or any(
        arch.lower().startswith("lfm2") for arch in architectures
    )


def get_tokenizer_config(
    model_name: str,
    trust_remote_code: bool = False,
) -> dict[str, Any]:
    """
    Get tokenizer configuration with model-specific fixes.

    This function centralizes tokenizer configuration to ensure consistent
    behavior across different modules.

    Args:
        model_name: The model name or path.
        trust_remote_code: Whether to trust remote code.

    Returns:
        Dictionary of tokenizer configuration options.
    """
    config: dict[str, Any] = {"trust_remote_code": trust_remote_code}

    # Apply Qwen3 fix if needed
    if is_qwen3_model(model_name):
        config["eos_token"] = "<|im_end|>"
        logger.debug("Qwen3 detected: setting eos_token to <|im_end|>")

    if _is_lfm2_text_lm(model_name):
        config.setdefault("tool_parser_type", "pythonic")
        logger.debug("LFM2 text LM detected: setting tool_parser_type to pythonic")

    return config


def apply_qwen3_fix(
    tokenizer_config: dict[str, Any],
    model_name: str,
) -> dict[str, Any]:
    """
    Apply Qwen3 tokenizer fix to an existing config.

    Qwen3 has a known issue where eos_token changed from <|im_end|> to
    <|endoftext|>, but the chat template still uses <|im_end|>. This
    function applies the fix if needed.

    Args:
        tokenizer_config: Existing tokenizer configuration dict.
        model_name: The model name or path.

    Returns:
        Updated tokenizer configuration with Qwen3 fix applied if needed.
    """
    if is_qwen3_model(model_name):
        tokenizer_config["eos_token"] = "<|im_end|>"
        logger.debug("Qwen3 detected: setting eos_token to <|im_end|>")

    return tokenizer_config