

# ===== FILE: omlx/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""
omlx: LLM inference server, optimized for your Mac

This package provides native Apple Silicon GPU acceleration using
Apple's MLX framework and mlx-lm for LLMs.

Features:
- Continuous batching via vLLM-style scheduler
- OpenAI-compatible API server
- Paged KV cache with prefix sharing
- Tiered cache (GPU + paged SSD offloading)
"""

from omlx._version import __version__

_LAZY = {
    "Request": "omlx.request",
    "RequestOutput": "omlx.request",
    "RequestStatus": "omlx.request",
    "SamplingParams": "omlx.request",
    "Scheduler": "omlx.scheduler",
    "SchedulerConfig": "omlx.scheduler",
    "SchedulerOutput": "omlx.scheduler",
    "EngineCore": "omlx.engine_core",
    "AsyncEngineCore": "omlx.engine_core",
    "EngineConfig": "omlx.engine_core",
    "BlockAwarePrefixCache": "omlx.cache.prefix_cache",
    "PagedCacheManager": "omlx.cache.paged_cache",
    "CacheBlock": "omlx.cache.paged_cache",
    "BlockTable": "omlx.cache.paged_cache",
    "PrefixCacheStats": "omlx.cache.stats",
    "PagedCacheStats": "omlx.cache.stats",
    "CacheStats": "omlx.cache.stats",
    "get_registry": "omlx.model_registry",
    "ModelOwnershipError": "omlx.model_registry",
}


def __getattr__(name: str):
    import importlib
    if name in _LAZY:
        mod = importlib.import_module(_LAZY[name])
        attr = "PagedCacheStats" if name == "CacheStats" else name
        val = getattr(mod, attr)
        globals()[name] = val
        return val
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

__all__ = [
    # Request management
    "Request",
    "RequestOutput",
    "RequestStatus",
    "SamplingParams",
    # Scheduler
    "Scheduler",
    "SchedulerConfig",
    "SchedulerOutput",
    # Engine
    "EngineCore",
    "AsyncEngineCore",
    "EngineConfig",
    # Model registry
    "get_registry",
    "ModelOwnershipError",
    # Prefix cache (paged SSD-only)
    "BlockAwarePrefixCache",
    # Paged cache (memory efficiency)
    "PagedCacheManager",
    "CacheBlock",
    "BlockTable",
    "PagedCacheStats",
    "CacheStats",  # Backward compatibility alias
    # Version
    "__version__",
]


# ===== FILE: omlx/_torch_stub.py =====

# SPDX-License-Identifier: Apache-2.0
"""Minimal ``torch`` stub for the DMG bundle.

xgrammar 0.2.3 declares ``torch>=1.10.0`` as a runtime dep, but oMLX never
exercises its torch-backed code paths: bitmasks are allocated as numpy
``int32`` buffers, the C++ binding fills them, and the MLX kernel applies the
mask. The torch dep is load-bearing only at *import time* — module-level code
in ``xgrammar.matcher``, ``xgrammar.testing``, ``xgrammar.contrib.hf`` and
``tvm_ffi.core`` does ``import torch`` plus a handful of attribute lookups.

Real torch is ~500 MB unpacked on macOS arm64 — too heavy to ship in the DMG.
This stub provides just enough of the torch surface for those modules to
finish loading. Code paths that would actually call into torch raise
``RuntimeError`` from the helpers below; oMLX never reaches them.

When a real torch is installed (pip / Homebrew flow) the stub is a no-op:
``install()`` checks ``importlib.util.find_spec('torch')`` first.
"""

from __future__ import annotations

import importlib.machinery
import importlib.metadata
import importlib.util
import logging
import os
import sys
import threading
import types

logger = logging.getLogger(__name__)

# xgrammar / tvm-ffi versions this stub is known to cover.
# This module is the *single source of truth* — packaging/build.py imports
# these constants to keep the DMG install pin in sync with the stub. Update
# both tuples here when bumping; the build script auto-tracks.
#
# Reachable-but-stubbed torch surface to be aware of when upgrading:
#   - ``torch.full``: ``xgrammar.allocate_token_bitmask`` calls it. oMLX
#     never invokes ``allocate_token_bitmask`` (we use the MLX kernel
#     path), but the symbol is re-exported from ``xgrammar.__init__``.
#     Any future caller that touches it will hit ``_unsupported("full")``
#     and surface a clear RuntimeError.
#   - ``torch.tensor`` returns a ``_StubTensor`` whose attribute access
#     raises a stub-identifying RuntimeError. Module-level
#     ``_FULL_MASK = torch.tensor(-1, ...)`` patterns succeed at import
#     time; any subsequent method call (.fill_, .item, ...) fails.
_TARGET_XGRAMMAR_VERSIONS = ("0.2.3",)
_TARGET_TVM_FFI_VERSIONS = ("0.1.11",)

# Serialize install() across threads. Without this, two threads that both
# pass the "torch" in sys.modules check race to build modules and overwrite
# each other's sys.modules['torch'] entry, leaving threads that already
# dereferenced the loser's module with stale references. Reachable today
# from concurrent HTTP handlers that call install() on first xgrammar use.
_INSTALL_LOCK = threading.Lock()
_INSTALLED = False


class _StubTensor:
    """Placeholder for ``torch.Tensor`` (annotations + isinstance checks).

    Any attribute access raises a clear RuntimeError so runtime use of a
    stubbed tensor (e.g. ``some_tensor.fill_(...)``) fails loudly with a
    pointer to the cause, rather than at the AttributeError level with a
    generic ``has no attribute 'fill_'`` message.
    """

    def __getattr__(self, name: str):
        # Let dunder probes (pickle, copy.deepcopy, descriptor lookups,
        # `hasattr` chains in third-party libs) fall through cleanly as
        # AttributeError — that's the documented `__getattr__` contract.
        # Real torch tensors lack many of these probed dunders anyway, so
        # raising AttributeError is the correct, distinguishable signal.
        if name.startswith("__") and name.endswith("__"):
            raise AttributeError(name)
        raise RuntimeError(
            f"_StubTensor.{name} is not implemented: oMLX ships a torch "
            "stub for xgrammar's import-time needs only. Reaching a real "
            "tensor method means a code path that needs real torch was "
            "exercised — install torch via pip/Homebrew or report this as "
            "a bug if the call originated inside oMLX."
        )


class _StubDtype:
    __slots__ = ("_name",)

    def __init__(self, name: str) -> None:
        self._name = name

    def __repr__(self) -> str:
        return f"torch.{self._name}"

    # Some xgrammar/tvm-ffi paths convert dtype to string via ``str(dt)``
    # rather than ``repr(dt)`` (e.g. ``to_cpp_dtype`` strips the "torch."
    # prefix). Match real torch's behaviour where ``str(torch.int32)`` is
    # ``"torch.int32"`` so those paths keep working.
    def __str__(self) -> str:
        return f"torch.{self._name}"


def _stub_tensor_factory(*args, **kwargs) -> _StubTensor:
    """torch.tensor(...) stub: returns a _StubTensor instance.

    Returning a real object (rather than None) means module-globals like
    xgrammar.matcher._FULL_MASK = torch.tensor(-1, dtype=...) succeed at
    import time. Any subsequent method call on the result (.fill_, .item,
    etc.) raises with a clear pointer via _StubTensor.__getattr__.
    """
    return _StubTensor()


def _false(*args, **kwargs) -> bool:
    return False


def _true(*args, **kwargs) -> bool:
    return True


def _zero(*args, **kwargs) -> int:
    return 0


def _unsupported(qualname: str):
    def _fn(*args, **kwargs):
        raise RuntimeError(
            f"torch.{qualname} is not available: this oMLX build ships a "
            "torch stub for xgrammar's import-time needs only. Install "
            "real torch via pip/Homebrew if you need this code path."
        )

    return _fn


# (canonical, alias) pairs — real torch aliases torch.int to torch.int32,
# torch.long to torch.int64, etc.; preserve those identities so code that
# does ``torch.int is torch.int32`` keeps working.
_DTYPE_ALIASES: tuple[tuple[str, tuple[str, ...]], ...] = (
    ("int32", ("int",)),
    ("int16", ("short",)),
    ("int64", ("long",)),
    ("float16", ("half",)),
    ("float32", ("float",)),
    ("float64", ("double",)),
    ("int8", ()),
    ("uint8", ()),
    ("bfloat16", ()),
    ("bool", ()),
)

_TENSOR_ALIASES = (
    "Tensor", "LongTensor", "FloatTensor", "IntTensor", "ByteTensor",
    "DoubleTensor", "HalfTensor", "BoolTensor", "ShortTensor",
)


# Names that xgrammar / tvm_ffi probe via getattr(torch, name) for
# feature-detection — they catch AttributeError and fall back gracefully.
# Logging WARNING for these floods the log on every model load (one per
# name per process) with diagnostics that aren't actually actionable.
# Demote known-probed names to DEBUG; everything else stays WARNING so
# genuinely-missing attributes surface in operator logs.
_KNOWN_PROBE_NAMES: frozenset[str] = frozenset({
    # Integer dtypes added post-torch-2.0 that tvm_ffi.dtypes enumerates
    "uint16", "uint32", "uint64",
    # FP8 / FP4 dtypes (probed by tvm_ffi.dtypes' dtype-mapping table)
    "float8_e4m3fn", "float8_e4m3fnuz",
    "float8_e5m2", "float8_e5m2fnuz",
    "float8_e8m0fnu",
    "float4_e2m1fn_x2",
})


def _make_top_level_torch_getattr() -> "callable":
    """Return a ``__getattr__`` for the stub's top-level torch module.

    Real-torch users who reach an unset attribute would get an
    ``AttributeError``; consumers that probe with ``hasattr`` rely on that.
    But we *also* want a clearly-identifiable message when downstream
    libraries (transformers, accelerate, etc.) reach for a torch surface
    we never stubbed — so this raises ``AttributeError`` whose message
    pinpoints the omlx stub. ``pkgutil.iter_modules(torch.__path__)`` and
    similar discovery paths see the empty ``__path__`` and short-circuit
    before hitting this.
    """

    _missing_attr_logged: set[str] = set()

    def __getattr__(name: str):  # noqa: N807
        # Surface the miss at WARNING level so a future xgrammar release
        # reaching for a new torch attribute is diagnosable from logs
        # before the AttributeError surfaces in a request handler. Rate-
        # limit per name so repeated probes (e.g. hasattr() under a
        # loop) don't flood the journal — once per name per process is
        # enough to identify the gap. Known-probed dtype names log at
        # DEBUG because xgrammar / tvm_ffi catch the AttributeError and
        # the WARNING is pure noise on every model load.
        if name not in _missing_attr_logged:
            _missing_attr_logged.add(name)
            level = logging.DEBUG if name in _KNOWN_PROBE_NAMES else logging.WARNING
            logger.log(
                level,
                "oMLX torch stub missing attribute: torch.%s "
                "(install real torch if this is load-bearing)",
                name,
            )
        # Dunder probes always fall through as AttributeError so pickling,
        # copy.deepcopy, and similar Python machinery work as expected.
        raise AttributeError(
            f"torch.{name!s} is not provided by the oMLX torch stub. "
            "Install real torch via pip/Homebrew if this attribute is "
            "actually needed."
        )

    return __getattr__


def _build_modules() -> dict[str, types.ModuleType]:
    torch = types.ModuleType("torch")
    for alias in _TENSOR_ALIASES:
        setattr(torch, alias, _StubTensor)
    torch.dtype = _StubDtype
    torch.__version__ = "0.0.0+omlx-stub"
    # Pin the stub as the source of truth for the xgrammar version it
    # targets; packaging/build.py imports this constant to stay in sync.
    # (Module-level constant lives at the top of this file.)
    for canonical, aliases in _DTYPE_ALIASES:
        dt = _StubDtype(canonical)
        setattr(torch, canonical, dt)
        for a in aliases:
            setattr(torch, a, dt)
    torch.tensor = _stub_tensor_factory
    torch.full = _unsupported("full")
    torch.zeros = _unsupported("zeros")
    torch.from_dlpack = _unsupported("from_dlpack")

    cuda = types.ModuleType("torch.cuda")
    cuda.is_available = _false
    cuda.device_count = _zero

    cuda_amp_common = types.ModuleType("torch.cuda.amp.common")
    cuda_amp_common.amp_definitely_not_available = _true
    cuda_amp = types.ModuleType("torch.cuda.amp")
    cuda_amp.common = cuda_amp_common
    cuda.amp = cuda_amp

    class _Stream:
        pass

    cuda.Stream = _Stream
    torch.cuda = cuda

    backends_mps = types.ModuleType("torch.backends.mps")
    backends_mps.is_available = _false
    backends_mps.is_built = _false
    backends_cudnn = types.ModuleType("torch.backends.cudnn")
    backends_cudnn.deterministic = False
    backends_cudnn.benchmark = False
    backends = types.ModuleType("torch.backends")
    backends.mps = backends_mps
    backends.cudnn = backends_cudnn
    torch.backends = backends

    version = types.ModuleType("torch.version")
    version.cuda = None
    version.hip = None
    torch.version = version

    nn_functional = types.ModuleType("torch.nn.functional")
    nn_functional.pad = _unsupported("nn.functional.pad")
    nn = types.ModuleType("torch.nn")
    nn.functional = nn_functional
    torch.nn = nn

    utils_dlpack = types.ModuleType("torch.utils.dlpack")
    utils_dlpack.to_dlpack = _unsupported("utils.dlpack.to_dlpack")
    utils = types.ModuleType("torch.utils")
    utils.dlpack = utils_dlpack
    torch.utils = utils

    # Top-level __getattr__ so a future xgrammar that reaches into a
    # torch surface we never stubbed (e.g. ``torch.compile``,
    # ``torch.distributed``) fails with a stub-identifying message rather
    # than a cryptic ``AttributeError: module 'torch' has no attribute…``.
    torch.__getattr__ = _make_top_level_torch_getattr()

    return {
        "torch": torch,
        "torch.cuda": cuda,
        "torch.cuda.amp": cuda_amp,
        "torch.cuda.amp.common": cuda_amp_common,
        "torch.backends": backends,
        "torch.backends.mps": backends_mps,
        "torch.backends.cudnn": backends_cudnn,
        "torch.version": version,
        "torch.nn": nn,
        "torch.nn.functional": nn_functional,
        "torch.utils": utils,
        "torch.utils.dlpack": utils_dlpack,
    }


def install() -> bool:
    """Install the stub into ``sys.modules`` if no real torch is available.

    Returns True if the stub was installed (or had been installed previously),
    False if a real torch was found and left alone.

    Thread-safe — concurrent callers (e.g. multiple FastAPI handlers hitting
    the xgrammar entry points in parallel) serialize on _INSTALL_LOCK.
    """
    global _INSTALLED
    needs_version_check = False
    with _INSTALL_LOCK:
        if _INSTALLED:
            return True

        if "torch" in sys.modules:
            already_stub = getattr(
                sys.modules["torch"], "__version__", ""
            ).endswith("+omlx-stub")
            _INSTALLED = already_stub
            return already_stub

        try:
            if importlib.util.find_spec("torch") is not None:
                # Real torch is on the path — leave it alone, install() is
                # a no-op. Don't mark _INSTALLED so a future sys.modules
                # reset (e.g. in tests) re-evaluates. Crucially, also DO
                # NOT touch ``TVM_FFI_DISABLE_TORCH_C_DLPACK`` — the user
                # has real torch and the tvm-ffi/torch-C-DLPack JIT path
                # may be their preferred fast path.
                return False
        except Exception:
            # find_spec can raise on broken parent packages, partial
            # installs, or weird import hooks. Treat as "no torch" — the
            # stub is the safe fallback.
            pass

        # No real torch — disable tvm_ffi's JIT torch-C-DLPack extension
        # before any tvm-ffi / xgrammar import. Without this,
        # tvm_ffi/_optional_torch_c_dlpack tries to JIT a C extension
        # against our stub at first import, spawns a doomed Python
        # subprocess that fails to ``import torch.utils.cpp_extension``
        # (the stub does not provide it), and surfaces a misleading
        # "Failed to JIT torch c dlpack extension" warning to users on
        # every cold start. The guard inside that module honours this
        # env var and skips the JIT path entirely.
        os.environ.setdefault("TVM_FFI_DISABLE_TORCH_C_DLPACK", "1")

        for name, mod in _build_modules().items():
            # ``__spec__`` must be a real ModuleSpec (not None) so that
            # ``importlib.util.find_spec`` succeeds when called by
            # transformers and other consumers. ``__version__`` is a
            # clearly-fake value so transformers refuses to take the
            # torch-modeling path.
            mod.__spec__ = importlib.machinery.ModuleSpec(name, loader=None)
            mod.__loader__ = None
            if "." not in name:
                mod.__path__ = []  # type: ignore[attr-defined]
            sys.modules[name] = mod
        _INSTALLED = True
        needs_version_check = True

    # Fire the version-drift check OUTSIDE the install lock — it reads
    # distribution metadata from disk and there is no reason to hold up
    # concurrent install() callers behind it. install() is idempotent at
    # this point — _INSTALLED is set and any racing caller short-circuits
    # at the top of the lock.
    if needs_version_check:
        try:
            warn_if_unexpected_versions()
        except Exception:  # pragma: no cover — defensive
            pass
    return True


def warn_if_unexpected_versions() -> None:
    """Log a warning when installed xgrammar / tvm-ffi versions drift past
    the versions this stub was tested against.

    Reads distribution metadata instead of module attributes: xgrammar
    exposes no ``__version__`` (checked on 0.2.3 and 0.2.4), so the old
    ``getattr(xgrammar, "__version__", None)`` probe never fired and the
    drift warning was dead code. Metadata also avoids importing the heavy
    C++ extension just to read a version string. Best-effort: silent when
    a distribution is not installed.
    """
    try:
        v = importlib.metadata.version("xgrammar")
        if v not in _TARGET_XGRAMMAR_VERSIONS:
            logger.warning(
                "xgrammar %s is not in the torch-stub target set %s; "
                "structured output may fail at runtime. Update the stub "
                "or pin xgrammar back.",
                v,
                _TARGET_XGRAMMAR_VERSIONS,
            )
    except Exception:
        pass
    try:
        v = importlib.metadata.version("apache-tvm-ffi")
        if v not in _TARGET_TVM_FFI_VERSIONS:
            logger.warning(
                "apache-tvm-ffi %s is not in the torch-stub target set %s; "
                "structured output may fail at runtime.",
                v,
                _TARGET_TVM_FFI_VERSIONS,
            )
    except Exception:
        pass


# ===== FILE: omlx/_version.py =====

__version__ = "0.5.4rc2"


# ===== FILE: omlx/adapter/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Adapters for different model formats.

This module provides adapters for parsing different model output formats
such as Harmony (gpt-oss).
"""

from .harmony import HarmonyStreamingParser

__all__ = ["HarmonyStreamingParser"]


# ===== FILE: omlx/adapter/gemma4.py =====

# SPDX-License-Identifier: Apache-2.0
"""Gemma 4 reasoning-channel output parsing and message extraction."""

from __future__ import annotations

import json
import re
from typing import Any

from ..api.utils import _PRESERVE_BOUNDARY_KEY
from ..utils.tokenizer import create_streaming_detokenizer
from .output_parser import OutputParserFinalizeResult, OutputParserTokenResult

_OPEN_MARKER = "<|channel>thought\n"
_OPEN_MARKER_BARE = "<|channel>"
_CLOSE_MARKER = "<channel|>"
_TURN_END_MARKER = "<turn|>"
_TOOL_RESPONSE_OPEN = "<|tool_response>"
_TOOL_RESPONSE_CLOSE = "<tool_response|>"
_THINK_OPEN = "<think>\n"
_THINK_CLOSE = "</think>\n"

_LEADING_THOUGHT_RE = re.compile(
    r"\A\s*(?:(?:<think>.*?</think>|<\|channel>.*?<channel\|>)\s*)+",
    re.DOTALL,
)

# Matches the STRAY bare-token spellings (<|tool_call> and <tool_call|>),
# not the template's well-formed closing form (</tool_call|> with slash).
_PROTOCOL_MARKER_RE = re.compile(r"<\|tool_call>|<tool_call\|>")


def _strip_protocol_markers(text: Any) -> Any:
    """Remove stray <|tool_call> / <tool_call|> tokens from assistant content."""
    if not isinstance(text, str) or not text:
        return text
    return _PROTOCOL_MARKER_RE.sub("", text)


def _try_parse_json(s: str) -> Any:
    """Parse string as JSON if possible, otherwise return as-is."""
    if not isinstance(s, str):
        return s
    s = s.strip()
    if not s or not (s.startswith("{") or s.startswith("[")):
        return s
    try:
        return json.loads(s)
    except (json.JSONDecodeError, ValueError):
        return s


def _strip_thinking(text: Any) -> Any:
    """Remove leading ``<think>...</think>`` or raw ``<|channel>...<channel|>`` spans.

    Gemma 4's multi-turn rule requires that only the final visible answer
    is kept in chat history. Clients such as Open WebUI replay the full
    assistant content (including the rendered ``<think>`` block, or the
    raw protocol form when a client preserves it). Feeding prior thought
    blocks back primes the model to emit malformed channel markers on the
    next turn, which then leak into user-facing output.

    The match is anchored to the start of the message: the rendered thought
    block always precedes the visible answer, so this catches every
    legitimate occurrence while leaving inline mentions (e.g. an assistant
    explaining how ``<think>`` tags work) untouched.
    """
    if not isinstance(text, str) or not text:
        return text
    return _LEADING_THOUGHT_RE.sub("", text, count=1)


def extract_gemma4_messages(
    messages: list[Any],
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    consolidate_system_messages: bool = True,
) -> list[dict]:
    """Convert OpenAI-format messages to Gemma 4 chat-template format.

    The Gemma 4 chat template does not handle ``role=tool`` messages.
    Tool results must instead appear on a model-role turn as a
    ``tool_responses`` list, where each entry is::

        {"name": "<function_name>", "response": <dict_or_scalar>}

    This function:
    - Passes non-tool messages through unchanged.
    - Preserves ``tool_calls`` on assistant turns (template renders them
      as ``<|tool_call>...</tool_call|>``).
    - Folds consecutive ``role=tool`` messages that follow an assistant
      turn into a single ``{"role": "assistant", "tool_responses": [...]}``
      message, resolving function names from the preceding tool_calls by
      ``tool_call_id``.  Falls back to the raw ``tool_call_id`` as the
      name when no match is found.
    - JSON-parses tool result content into a dict/list where possible so
      the template renders structured responses correctly.

    Args:
        messages: OpenAI-format Message objects or dicts.
        max_tool_result_tokens: Maximum token count for tool results
            (truncation applied when tokenizer is provided).
        tokenizer: Tokenizer for optional truncation.
        consolidate_system_messages: When True, preserve the legacy behavior
            of moving all system/developer messages to the leading system
            prompt. Server routes pass False and defer that decision until the
            model chat template can be probed.

    Returns:
        List of dicts ready for ``tokenizer.apply_chat_template``.
    """
    from ..api.utils import (
        _extract_text_from_content_list,
    )  # avoid circular at module level

    processed: list[dict] = []

    # Build index of message objects as plain dicts
    raw: list[dict] = []
    for msg in messages:
        if hasattr(msg, "model_dump"):
            raw.append(msg.model_dump())
        elif isinstance(msg, dict):
            raw.append(dict(msg))
        else:
            raw.append(
                {
                    "role": getattr(msg, "role", "user"),
                    "content": getattr(msg, "content", ""),
                }
            )

    i = 0
    while i < len(raw):
        msg = raw[i]
        role = msg.get("role", "user")

        if role == "developer":
            role = "system"

        if role == "tool":
            # Orphaned tool result with no preceding assistant turn — attach
            # to a synthetic assistant turn with no content.
            tool_call_id = msg.get("tool_call_id", "")
            content = msg.get("content", "")
            if isinstance(content, list):
                content = _extract_text_from_content_list(content)
            if max_tool_result_tokens and tokenizer and content:
                from ..api.anthropic_utils import truncate_tool_result

                content = truncate_tool_result(
                    content, max_tool_result_tokens, tokenizer
                )
            response = _try_parse_json(content)
            # Fallback name: use tool_call_id
            processed.append(
                {
                    "role": "assistant",
                    "content": "",
                    "tool_responses": [
                        {"name": tool_call_id or "unknown", "response": response}
                    ],
                    _PRESERVE_BOUNDARY_KEY: True,
                }
            )
            i += 1
            continue

        if role == "assistant":
            # Build a tool_call_id → function_name lookup from this turn's calls.
            tc_id_to_name: dict[str, str] = {}
            tool_calls_raw = msg.get("tool_calls") or []
            for tc in tool_calls_raw:
                if isinstance(tc, dict):
                    tc_id = tc.get("id", "")
                    func_name = (tc.get("function") or {}).get("name", "")
                else:
                    tc_id = getattr(tc, "id", "")
                    func = getattr(tc, "function", None)
                    func_name = getattr(func, "name", "") if func else ""
                if tc_id:
                    tc_id_to_name[tc_id] = func_name

            # Extract content
            content = msg.get("content", "")
            if isinstance(content, list):
                content = _extract_text_from_content_list(content)
            # Per Gemma 4's multi-turn rule, prior thought blocks must not
            # be fed back into the next turn. Strip them before rendering.
            content = _strip_thinking(content)
            content = _strip_protocol_markers(content)

            out_msg: dict = {"role": "assistant", "content": content or ""}

            # Preserve tool_calls for template rendering
            if tool_calls_raw:
                out_calls = []
                for tc in tool_calls_raw:
                    if isinstance(tc, dict):
                        func = tc.get("function") or {}
                        out_calls.append(
                            {
                                "id": tc.get("id", ""),
                                "function": {
                                    "name": func.get("name", ""),
                                    "arguments": _try_parse_json(
                                        func.get("arguments", "{}")
                                    ),
                                },
                            }
                        )
                    else:
                        func = getattr(tc, "function", None)
                        args_str = getattr(func, "arguments", "{}") if func else "{}"
                        out_calls.append(
                            {
                                "id": getattr(tc, "id", ""),
                                "function": {
                                    "name": getattr(func, "name", "") if func else "",
                                    "arguments": _try_parse_json(args_str),
                                },
                            }
                        )
                out_msg["tool_calls"] = out_calls
                out_msg[_PRESERVE_BOUNDARY_KEY] = True

            processed.append(out_msg)
            i += 1

            # Consume any immediately following tool results into a
            # single model turn with tool_responses.
            tool_responses = []
            while i < len(raw) and raw[i].get("role") == "tool":
                tr = raw[i]
                tc_id = tr.get("tool_call_id", "")
                tr_content = tr.get("content", "")
                if isinstance(tr_content, list):
                    tr_content = _extract_text_from_content_list(tr_content)
                if max_tool_result_tokens and tokenizer and tr_content:
                    from ..api.anthropic_utils import truncate_tool_result

                    tr_content = truncate_tool_result(
                        tr_content, max_tool_result_tokens, tokenizer
                    )
                response = _try_parse_json(tr_content)
                name = tc_id_to_name.get(tc_id) or tc_id or "unknown"
                tool_responses.append({"name": name, "response": response})
                i += 1

            if tool_responses:
                # Attach tool_responses to the SAME assistant message that
                # has tool_calls.  The Gemma 4 chat template checks for
                # tool_responses on the current message (lines 261-267)
                # BEFORE falling back to a forward-scan for role='tool'
                # messages (lines 268-302).  Putting them on a separate
                # assistant message causes both paths to miss, producing a
                # corrupt bare <|tool_response> tag and making the model
                # loop on the same tool call.
                out_msg["tool_responses"] = tool_responses
            continue

        # All other roles (user, system)
        # Preserve image_url and input_audio parts for VLM processing
        content = msg.get("content", "")
        if isinstance(content, list):
            from ..api.utils import _extract_multimodal_content_list

            multimodal_parts = _extract_multimodal_content_list(content)
            multimodal_types = {"image_url", "input_audio"}
            has_multimodal = any(
                p.get("type") in multimodal_types for p in multimodal_parts
            )
            if has_multimodal:
                content = multimodal_parts
            else:
                content = _extract_text_from_content_list(content)
        out: dict = {"role": role, "content": content if content is not None else ""}
        processed.append(out)
        i += 1

    # Standard cleanup passes shared with other extractors
    from ..api.utils import (
        _consolidate_system_messages,
        _drop_void_assistant_messages,
        _merge_consecutive_roles,
    )

    cleaned = processed
    if consolidate_system_messages:
        cleaned = _consolidate_system_messages(cleaned)
    cleaned = _drop_void_assistant_messages(cleaned)
    return _merge_consecutive_roles(cleaned)


def _matching_prefix_len(text: str, marker: str) -> int:
    """Return longest suffix of ``text`` that is a prefix of ``marker``."""
    max_len = min(len(text), len(marker) - 1)
    for size in range(max_len, 0, -1):
        if text.endswith(marker[:size]):
            return size
    return 0


class Gemma4OutputParserSession:
    """Suppress Gemma 4 protocol markers and re-emit thought blocks as ``<think>`` tags."""

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        self._tokenizer = tokenizer
        self._buffer = ""
        self._in_thought = False
        self._text_mode = False

        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()

    def _append_text(
        self,
        stream_parts: list[str],
        visible_parts: list[str],
        text: str,
    ) -> None:
        if not text:
            return
        stream_parts.append(text)
        visible_parts.append(text)

    def _active_markers(self) -> list[str]:
        # Channel open/close are tracked unconditionally so a stray
        # ``<channel|>`` outside a thought block (occasionally emitted in long
        # multi-turn contexts) is absorbed instead of leaking into visible
        # text. ``_OPEN_MARKER_BARE`` is a defensive fallback for malformed
        # opens (e.g. ``<|channel>thought<channel|>`` with no newline, or a
        # bare ``<|channel>`` emitted when the model is confused by polluted
        # history). Tool-call markup is intentionally not tracked here — the
        # downstream ``ToolCallStreamFilter`` removes it from stream deltas
        # while ``parse_tool_calls`` still sees the raw markers in
        # ``output_text`` for extraction.
        return [
            _OPEN_MARKER,
            _OPEN_MARKER_BARE,
            _CLOSE_MARKER,
            _TURN_END_MARKER,
            _TOOL_RESPONSE_OPEN,
            _TOOL_RESPONSE_CLOSE,
        ]

    @staticmethod
    def _find_next_marker(
        source: str, pos: int, markers: list[str]
    ) -> tuple[int, str] | tuple[None, None]:
        next_idx: int | None = None
        next_marker: str | None = None
        for marker in markers:
            idx = source.find(marker, pos)
            if idx == -1:
                continue
            if next_idx is None or idx < next_idx:
                next_idx = idx
                next_marker = marker
        return next_idx, next_marker

    def _consume_text(
        self, text: str, *, final: bool = False
    ) -> OutputParserTokenResult:
        source = self._buffer + text
        self._buffer = ""

        stream_parts: list[str] = []
        visible_parts: list[str] = []
        pos = 0

        while pos < len(source):
            markers = self._active_markers()
            idx, marker = self._find_next_marker(source, pos, markers)

            if idx is None or marker is None:
                remainder = source[pos:]
                if not final:
                    keep = max(
                        _matching_prefix_len(remainder, marker_text)
                        for marker_text in markers
                    )
                    if keep:
                        emit = remainder[:-keep]
                        self._buffer = remainder[-keep:]
                    else:
                        emit = remainder
                else:
                    emit = remainder

                self._append_text(stream_parts, visible_parts, emit)
                break

            # Streaming defer: a bare ``<|channel>`` (or ``<|channel>thought``
            # without trailing newline) at the end of the source could still
            # extend to the canonical ``<|channel>thought\n`` once more
            # tokens arrive. Buffer and wait so the canonical match wins.
            if not final and marker == _OPEN_MARKER_BARE:
                suffix = source[idx:]
                if len(suffix) < len(_OPEN_MARKER) and _OPEN_MARKER.startswith(suffix):
                    self._append_text(stream_parts, visible_parts, source[pos:idx])
                    self._buffer = suffix
                    return OutputParserTokenResult(
                        stream_text="".join(stream_parts),
                        visible_text="".join(visible_parts),
                    )

            self._append_text(stream_parts, visible_parts, source[pos:idx])

            advance = len(marker)

            if marker == _OPEN_MARKER:
                # Nested open while already in a thought block: drop the stray
                # marker without re-emitting ``<think>`` to keep the structure
                # well-formed.
                if not self._in_thought:
                    stream_parts.append(_THINK_OPEN)
                    visible_parts.append(_THINK_OPEN)
                    self._in_thought = True
            elif marker == _OPEN_MARKER_BARE:
                # Defensive fallback for malformed opens: ``<|channel>thought``
                # without the trailing newline, or a bare ``<|channel>`` with
                # an unrecognised channel name. Treat as a thought open and
                # absorb the optional ``thought`` keyword and newline so they
                # don't leak as visible text.
                if not self._in_thought:
                    stream_parts.append(_THINK_OPEN)
                    visible_parts.append(_THINK_OPEN)
                    self._in_thought = True
                after = idx + advance
                if source.startswith("thought\n", after):
                    advance += len("thought\n")
                elif source.startswith("thought", after):
                    advance += len("thought")
            elif marker == _CLOSE_MARKER:
                # Stray close outside a thought block: drop silently to keep
                # the marker out of visible content.
                if self._in_thought:
                    stream_parts.append(_THINK_CLOSE)
                    visible_parts.append(_THINK_CLOSE)
                    self._in_thought = False
            # _TURN_END_MARKER, _TOOL_RESPONSE_OPEN / _CLOSE: silent drop.

            pos = idx + advance

        return OutputParserTokenResult(
            stream_text="".join(stream_parts),
            visible_text="".join(visible_parts),
        )

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        if self._detokenizer is not None:
            self._detokenizer.add_token(token_id)
            text = self._detokenizer.last_segment
        else:
            text = self._tokenizer.decode([token_id])
        return self._consume_text(text)

    def process_text(self, text: str) -> OutputParserTokenResult:
        """Process an already-detokenized text segment.

        Engines that emit text segments instead of token ids (the serial
        diffusion lane detokenizes inside ``stream_diffusion_generate``)
        feed their output through this entry point so protocol markers
        are handled identically to the token-id path.  Switches the
        session to text mode so ``finalize`` does not flush the unused
        token detokenizer.
        """
        self._text_mode = True
        if not text:
            return OutputParserTokenResult(stream_text="", visible_text="")
        return self._consume_text(text)

    def finalize(self) -> OutputParserFinalizeResult:
        text = ""
        if self._detokenizer is not None and not self._text_mode:
            self._detokenizer.finalize()
            text = self._detokenizer.last_segment

        token_result = self._consume_text(text, final=True)

        stream_text = token_result.stream_text
        visible_text = token_result.visible_text

        if self._buffer:
            stream_text += self._buffer
            visible_text += self._buffer
            self._buffer = ""

        if self._in_thought:
            stream_text += _THINK_CLOSE
            visible_text += _THINK_CLOSE
            self._in_thought = False

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
        )


# ===== FILE: omlx/adapter/harmony.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Harmony format streaming parser for gpt-oss models.

Uses the official openai-harmony package for robust parsing.

Harmony protocol uses special tokens to structure messages:
- <|start|>: Begin message header
- <|channel|>: Mark channel type
- <|message|>: Transition to content
- <|end|>: End message
- <|return|>: Model completion signal
- <|call|>: Tool invocation signal

Message structure: <|start|>{header}<|channel|>{channel_name}<|message|>{content}<|end|>

Channels:
- final: User-visible response (plain text)
- analysis: Chain-of-thought reasoning (wrapped in <think>...</think> for streaming)
- commentary: Tool/function calls (non-streaming only)

gpt-oss also emits tool calls on the analysis channel with an explicit
``to=functions.*`` recipient; those are honored as tool calls when the
arguments form a JSON object (#2216).
"""

import json
import logging
import re
import time
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any

from openai_harmony import (
    HarmonyEncoding,
    Role,
    StreamableParser,
    load_harmony_encoding,
)

logger = logging.getLogger(__name__)

# Pattern to match <think>...</think> blocks
_THINK_TAG_PATTERN = re.compile(r"<think>.*?</think>\s*", re.DOTALL)

# Pre-allocated constants
_THINK_OPEN = "<think>\n"
_THINK_CLOSE = "</think>\n"

# Harmony special tokens that should not be streamed
_HARMONY_SPECIAL_TOKENS = [
    "<|start|>",
    "<|end|>",
    "<|message|>",
    "<|channel|>",
    "<|return|>",
    "<|call|>",
    "<|constrain|>",
]


def _has_no_real_recipient(recipient: str | None) -> bool:
    """Return True when the parser only preserved the primed assistant header."""
    return recipient is None or recipient == "<|start|>assistant"


def _message_content_text(msg: Any) -> str:
    """Concatenate the text parts of a parsed Harmony message."""
    text = ""
    content = getattr(msg, "content", None)
    if content is not None:
        for part in content:
            part_text = getattr(part, "text", None)
            if isinstance(part_text, str):
                text += part_text
    return text


def _is_tool_call_message(msg: Any) -> bool:
    """Return True when a parsed message is a genuine tool invocation.

    Commentary is the canonical tool-call channel. gpt-oss also emits tool
    calls on the analysis channel with an explicit functions.* recipient
    (#2216); those count only when the arguments form a JSON object, so
    recipient-less tool syntax reasoned about in thinking and prose
    addressed to a tool both stay in the reasoning output (#2032).
    """
    recipient = getattr(msg, "recipient", None)
    if not (isinstance(recipient, str) and recipient.startswith("functions.")):
        return False
    channel = getattr(msg, "channel", None)
    if channel == "commentary":
        return True
    if channel != "analysis":
        return False
    try:
        return isinstance(json.loads(_message_content_text(msg)), dict)
    except ValueError:
        return False


@lru_cache(maxsize=1)
def load_harmony_gpt_oss_encoding() -> HarmonyEncoding:
    """Load the Harmony gpt-oss encoding with a small retry window."""
    last_error: Exception | None = None
    for attempt in range(3):
        try:
            return load_harmony_encoding("HarmonyGptOss")
        except Exception as exc:
            last_error = exc
            if attempt == 2:
                break
            delay = 0.5 * (2**attempt)
            logger.warning(
                "Failed to load HarmonyGptOss encoding "
                "(attempt %d/3): %s; retrying in %.1fs",
                attempt + 1,
                exc,
                delay,
            )
            time.sleep(delay)

    assert last_error is not None
    raise last_error


def preprocess_harmony_messages(
    messages: list[dict[str, Any]]
) -> list[dict[str, Any]]:
    """
    Preprocess messages for Harmony (gpt-oss) models.

    - Strips <think> tags from assistant messages
    - Keeps tool role messages unchanged (chat_template handles conversion)

    The chat_template expects standard OpenAI format:
    - role: "tool" with tool_call_id and content
    - It uses last_tool_call.name from the previous assistant message
    - Generates: <|start|>functions.{name} to=assistant<|channel|>commentary<|message|>{content|tojson}<|end|>

    Args:
        messages: List of message dicts with 'role' and 'content' keys

    Returns:
        Messages preprocessed for Harmony format
    """
    if not messages:
        return []

    result = []

    for msg in messages:
        # Validate message is a dict
        if not isinstance(msg, dict):
            logger.warning(f"Skipping non-dict message: {type(msg)}")
            continue

        role = msg.get("role")

        if role == "assistant":
            content = msg.get("content", "")
            # Ensure content is a string (could be list in some formats)
            if isinstance(content, str):
                # Strip <think> tags
                if content and "<think>" in content:
                    content = _THINK_TAG_PATTERN.sub("", content).strip()
                    msg = {**msg, "content": content}
            elif content is not None:
                # Non-string content (e.g., list) - log but don't modify
                logger.debug(f"Assistant message has non-string content: {type(content)}")

            result.append(msg)

        else:
            # Pass through all other messages (user, tool, system, etc.) unchanged
            # Chat template handles tool messages directly using last_tool_call.name
            result.append(msg)

    return result


def _get_special_token_ids(tokenizer: Any) -> set[int]:
    """
    Get special token IDs from model tokenizer.

    Args:
        tokenizer: The model's tokenizer

    Returns:
        Set of special token IDs
    """
    special_ids = set()
    for token in _HARMONY_SPECIAL_TOKENS:
        try:
            token_id = tokenizer.convert_tokens_to_ids(token)
            if isinstance(token_id, int) and token_id >= 0:
                special_ids.add(token_id)
            else:
                logger.debug(f"Harmony special token '{token}' not found in tokenizer")
        except Exception as e:
            logger.debug(f"Failed to get ID for Harmony token '{token}': {e}")
    return special_ids


@dataclass
class HarmonyStreamingParser:
    """
    Streaming parser for Harmony format using official openai-harmony package.

    Parses tokens incrementally and routes them to appropriate channels.
    Returns token IDs instead of decoded text to allow proper UTF-8 handling
    via streaming detokenizer in the caller.

    Output routing:
    - analysis channel -> stream only (wrapped in <think>...</think>)
    - final channel -> stream and visible (stored in output_text)
    - commentary channel -> buffered for tool calls (non-streaming)

    The parser returns:
    - control_text: Control strings like <think>, </think>
    - stream_token: Token ID to stream (None if not streaming)
    - visible_token: Token ID to store (None if not storing)
    - is_stop: Whether this is a stop signal
    """

    tokenizer: Any

    # Internal state (initialized in __post_init__)
    _encoding: HarmonyEncoding = field(init=False, repr=False)
    _parser: StreamableParser = field(init=False, repr=False)
    _stop_tokens: set[int] = field(init=False, default_factory=set)
    _special_tokens: set[int] = field(init=False, default_factory=set)

    # <think> tag state
    _in_think_tag: bool = field(init=False, default=False)
    _prev_channel: str | None = field(init=False, default=None)

    # Passthrough mode: activated when streaming parser encounters an
    # unrecoverable error.  Tokens are still accumulated by the scheduler
    # (request.append_output_token) so parse_tool_calls_from_tokens can
    # extract tool calls at finalization.
    _passthrough_mode: bool = field(init=False, default=False)

    def __post_init__(self):
        """Initialize the official Harmony parser."""
        self._encoding = load_harmony_gpt_oss_encoding()
        # role=None allows the parser to handle tool-call headers
        # (e.g. "assistant to=functions.Write") which Role.ASSISTANT rejects.
        self._parser = StreamableParser(self._encoding, None, strict=False)
        self._stop_tokens = set(self._encoding.stop_tokens_for_assistant_actions())
        self._special_tokens = _get_special_token_ids(self.tokenizer)

        # Prime the parser with "<|start|>assistant" tokens.  The chat
        # template already includes these in the prompt, so the model's
        # first output token is <|channel|>, not <|start|>.  Without
        # priming, the parser rejects <|channel|> as unexpected.
        self._prime_parser(self._parser)

        logger.info(
            f"Harmony parser initialized: {len(self._special_tokens)} special tokens, "
            f"{len(self._stop_tokens)} stop tokens"
        )

    def _prime_parser(self, parser: StreamableParser) -> None:
        """Feed '<|start|>assistant' header tokens so parser expects <|channel|> next."""
        start_tokens = self._encoding.encode(
            "<|start|>assistant", allowed_special="all"
        )
        for t in start_tokens:
            parser.process(t)

    def process_token(
        self, token_id: int
    ) -> tuple[str, int | None, int | None, bool]:
        """
        Process a single token and return routing information.

        This method routes tokens to appropriate channels without decoding.
        The caller should use streaming detokenizer to decode the returned
        token IDs for proper UTF-8 handling.

        Args:
            token_id: The token ID to process.

        Returns:
            Tuple of:
            - control_text: Control strings (<think>, </think>, etc.)
            - stream_token: Token ID to stream (None to skip)
            - visible_token: Token ID to store in output_text (None to skip)
            - is_stop: True if this is a stop token
        """
        # Check if this is a special token (should not be streamed)
        is_special_token = token_id in self._special_tokens
        is_stop = token_id in self._stop_tokens
        was_analysis = self._prev_channel == "analysis"

        # Passthrough: parser crashed earlier, buffer all tokens silently.
        # Tokens are still tracked by the scheduler for non-streaming tool
        # call extraction at finalization.
        if self._passthrough_mode:
            return "", None, None, is_stop

        try:
            self._parser.process(token_id)
        except Exception as e:
            logger.warning(
                f"Harmony streaming parser error, switching to passthrough: {e}"
            )
            self._passthrough_mode = True
            control_text = ""
            if self._in_think_tag:
                control_text = _THINK_CLOSE
                self._in_think_tag = False
            return control_text, None, None, is_stop

        channel = self._parser.current_channel
        control_text = ""

        # Harmony uses the same end token for analysis, final, and tool/action
        # messages. Ending analysis should let generation continue into the
        # final channel; ending other channels should stop the request.
        if was_analysis and is_stop:
            is_stop = False

        # Handle channel transitions for <think> tags
        if channel != self._prev_channel:
            # Close previous analysis channel
            if self._in_think_tag and self._prev_channel == "analysis":
                control_text += _THINK_CLOSE
                self._in_think_tag = False
            # Open new analysis channel
            if channel == "analysis" and not self._in_think_tag:
                control_text += _THINK_OPEN
                self._in_think_tag = True
            self._prev_channel = channel

        # Special tokens should never be streamed or stored
        if is_special_token:
            return control_text, None, None, is_stop

        # Route based on channel
        if channel == "final":
            # final: stream AND store (same token for both)
            return control_text, token_id, token_id, is_stop
        elif channel == "analysis":
            # analysis: stream only (wrapped in <think>)
            return control_text, token_id, None, is_stop
        elif channel is None:
            # Channel not yet determined (still in header parsing)
            # Buffer token but don't stream
            return control_text, None, None, is_stop
        elif channel != "commentary" and _has_no_real_recipient(
            self._parser.current_recipient
        ):
            # Some fine-tunes emit misspelled channels (for example "mardown").
            # If there is no recipient, preserve the text as user-visible output.
            return control_text, token_id, token_id, is_stop
        else:
            # commentary etc: buffer only (for tool calls)
            return control_text, None, None, is_stop

    def get_stop_token_ids(self) -> set[int]:
        """Get Harmony stop token IDs."""
        return self._stop_tokens

    def get_tool_calls(self) -> list[dict[str, str]]:
        """Get accumulated tool calls from parsed messages."""
        tool_calls = []
        try:
            messages = self._parser.messages
            if not messages:
                return tool_calls

            for msg in messages:
                if not _is_tool_call_message(msg):
                    continue

                name = msg.recipient[10:]  # Remove "functions." prefix
                content = _message_content_text(msg)

                tool_calls.append({"name": name, "arguments": content})
                logger.info(f"Extracted tool call: {name}, arguments={content}")

        except Exception as e:
            logger.warning(f"Error extracting tool calls: {e}")

        return tool_calls

    def finalize(self) -> str:
        """
        Finalize parsing and close any open tags.

        Returns:
            Any remaining control text (e.g., closing </think> tag).
        """
        try:
            self._parser.process_eos()
        except Exception as e:
            # Can fail if message is incomplete (e.g., missing <|end|>)
            # This is expected in some cases, so just log and continue
            logger.debug(f"Harmony parser process_eos failed (expected for incomplete messages): {e}")

        if self._in_think_tag:
            self._in_think_tag = False
            return _THINK_CLOSE
        return ""

    def reset(self) -> None:
        """Reset parser state for a new request."""
        self._parser = StreamableParser(self._encoding, None, strict=False)
        self._prime_parser(self._parser)
        self._in_think_tag = False
        self._prev_channel = None
        self._passthrough_mode = False

    @property
    def current_channel(self) -> str | None:
        """Get current channel name."""
        return self._parser.current_channel

    @property
    def current_recipient(self) -> str | None:
        """Get current recipient (for tool calls)."""
        return self._parser.current_recipient


def parse_tool_calls_from_tokens(
    token_ids: list[int],
    prepend_start: bool = True,
) -> tuple[str, str, list[dict[str, str]]]:
    """
    Parse a complete Harmony token sequence (non-streaming).

    Args:
        token_ids: Model output token ID list
        prepend_start: Whether to prepend "<|start|>assistant" tokens.
            Set to False if token_ids already includes start tokens.

    Returns:
        (output_text, analysis_text, tool_calls)
        - output_text: Text from the final channel
        - analysis_text: Chain-of-thought text from the analysis channel
        - tool_calls: [{"name": "...", "arguments": "..."}]
    """
    if not token_ids:
        return "", "", []

    try:
        encoding = load_harmony_gpt_oss_encoding()

        start_tokens = encoding.encode("<|start|>assistant", allowed_special="all")
        has_start = list(token_ids[: len(start_tokens)]) == start_tokens

        # The normal chat template includes "<|start|>assistant" in the prompt,
        # so completions start from "<|channel|>" and need the prefix restored.
        # Budget-forced or recovered Harmony completions may already include it.
        if prepend_start and not has_start:
            full_token_ids = start_tokens + list(token_ids)
        else:
            full_token_ids = list(token_ids)

        # Decode tokens for debugging
        decoded_text = encoding.decode(full_token_ids)
        logger.info(f"parse_tool_calls input ({len(full_token_ids)} tokens): {decoded_text[:300]}...")

        messages = encoding.parse_messages_from_completion_tokens(
            full_token_ids,
            role=Role.ASSISTANT,
            strict=False,
        )

        logger.info(f"Parsed {len(messages)} messages")
        for i, msg in enumerate(messages):
            content_count = len(msg.content) if msg.content else 0
            logger.info(
                f"Message {i}: channel={msg.channel}, recipient={msg.recipient}, "
                f"content_count={content_count}"
            )

        output_text = ""
        analysis_text = ""
        tool_calls = []

        for msg in messages:
            # Safely get content
            msg_content = getattr(msg, "content", None)
            if msg_content is None:
                continue

            # Checked before the analysis branch so analysis-channel tool
            # calls do not leak their arguments into reasoning text.
            if _is_tool_call_message(msg):
                name = msg.recipient[10:]  # Remove "functions." prefix
                tool_calls.append(
                    {"name": name, "arguments": _message_content_text(msg)}
                )

            elif msg.channel == "final":
                # Extract text from final channel
                for content in msg_content:
                    text = getattr(content, "text", None)
                    if isinstance(text, str):
                        output_text += text

            elif msg.channel == "analysis":
                # Extract chain-of-thought text from analysis channel
                for content in msg_content:
                    text = getattr(content, "text", None)
                    if isinstance(text, str):
                        analysis_text += text

            elif msg.channel != "commentary" and _has_no_real_recipient(msg.recipient):
                # Preserve malformed/unknown assistant channels as visible text
                # instead of returning an empty assistant message.
                for content in msg_content:
                    text = getattr(content, "text", None)
                    if isinstance(text, str):
                        output_text += text

        return output_text, analysis_text, tool_calls

    except Exception as e:
        logger.warning(f"Error parsing tool calls from tokens: {e}")
        return "", "", []


# ===== FILE: omlx/adapter/output_parser.py =====

# SPDX-License-Identifier: Apache-2.0
"""Generic streamed output parser sessions.

This module provides a tiny scheduler-facing abstraction for protocol-specific
output parsing.  A parser session owns any protocol state needed while a single
request is generating (e.g. Harmony channel parsing or Gemma 4 reasoning marker
suppression) and exposes a uniform token-by-token interface.
"""

from __future__ import annotations

import json
import logging
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Protocol

from ..utils.tokenizer import (
    create_streaming_detokenizer,
    is_gemma4_model,
    is_harmony_model,
)
from .harmony import HarmonyStreamingParser, parse_tool_calls_from_tokens

logger = logging.getLogger(__name__)


@dataclass
class OutputParserTokenResult:
    """Per-token parser result returned during streaming."""

    stream_text: str = ""
    visible_text: str = ""
    is_stop: bool = False
    record_token: bool | None = None


@dataclass
class OutputParserFinalizeResult:
    """Final parser result returned once a request finishes."""

    stream_text: str = ""
    visible_text: str = ""
    output_text_prefix: str = ""
    tool_calls: list[dict[str, str]] = field(default_factory=list)
    finish_reason: str | None = None


class OutputParserSession(Protocol):
    """Protocol implemented by per-request output parser sessions."""

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        """Process one generated token."""

    def finalize(self) -> OutputParserFinalizeResult:
        """Flush any buffered output when generation ends."""


@dataclass(frozen=True)
class OutputParserFactory:
    """Factory for creating per-request parser sessions."""

    kind: str
    create_session: Callable[[Any], OutputParserSession]
    stop_token_ids: set[int] = field(default_factory=set)
    thinking_start_text: str | None = None
    thinking_start_output_text: str | None = None
    thinking_end_text: str | None = None
    thinking_end_trailing_text: str | None = None
    # Marker strings that must survive special-token stripping so the
    # parser session can see them in the text stream.  Engines that strip
    # special tokens during detokenization (e.g. the serial diffusion
    # lane) preserve the token ids of these markers and let the parser
    # session remove them instead.
    protocol_marker_texts: tuple[str, ...] = ()


class HarmonyOutputParserSession:
    """Scheduler-facing wrapper around ``HarmonyStreamingParser``."""

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        self._tokenizer = tokenizer
        self._parser = HarmonyStreamingParser(tokenizer)
        self._raw_token_ids: list[int] = []

        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        control_text, stream_token, visible_token, is_stop = self._parser.process_token(
            token_id
        )
        self._raw_token_ids.append(token_id)

        stream_text = control_text
        visible_text = ""

        if stream_token is not None:
            if self._detokenizer is not None:
                self._detokenizer.add_token(stream_token)
                decoded_text = self._detokenizer.last_segment
            else:
                decoded_text = self._tokenizer.decode([stream_token])

            stream_text += decoded_text
            if visible_token is not None:
                visible_text += decoded_text
        elif visible_token is not None:
            if self._detokenizer is not None:
                self._detokenizer.add_token(visible_token)
                visible_text += self._detokenizer.last_segment
            else:
                visible_text += self._tokenizer.decode([visible_token])

        return OutputParserTokenResult(
            stream_text=stream_text,
            visible_text=visible_text,
            is_stop=is_stop,
            record_token=True,
        )

    def finalize(self) -> OutputParserFinalizeResult:
        stream_text = self._parser.finalize()
        visible_text = ""

        if self._detokenizer is not None:
            self._detokenizer.finalize()
            final_text = self._detokenizer.last_segment
            if final_text:
                stream_text += final_text
                if self._parser.current_channel == "final":
                    visible_text += final_text

        _, analysis_text, tool_calls = parse_tool_calls_from_tokens(self._raw_token_ids)
        finish_reason = "tool_calls" if tool_calls else None

        output_text_prefix = (
            f"<think>\n{analysis_text}\n</think>\n" if analysis_text else ""
        )

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
            output_text_prefix=output_text_prefix,
            tool_calls=tool_calls,
            finish_reason=finish_reason,
        )


def _is_cohere2_moe_model(
    model_name: str,
    model_config: dict[str, Any] | None = None,
) -> bool:
    return model_config is not None and model_config.get("model_type") == "cohere2_moe"


_MINIMAX_M3_MODEL_TYPES = {"minimax_m3", "minimax_m3_vl"}
_MINIMAX_THINK_START = "<mm:think>"
_MINIMAX_THINK_END = "</mm:think>"
_MINIMAX_EOS_TOKEN = "[e~["
_MINIMAX_SPECIAL_TOKENS = (_MINIMAX_EOS_TOKEN, "]~b]", "]~!b[", "]!p~[", "]!d~[")
_MINIMAX_TOOL_CALL_START = "]<]minimax[>[<tool_call>"
_MINIMAX_TOOL_CALL_END = "]<]minimax[>[</tool_call>"
_DEEPSEEK_V4_TOOL_CALL_START = "<｜DSML｜tool_calls>"
_DEEPSEEK_V4_TOOL_CALL_END = "</｜DSML｜tool_calls>"


def _is_deepseek_v4_model(
    model_name: str,
    tokenizer: Any,
    model_config: dict[str, Any] | None = None,
) -> bool:
    model_type = str(model_config.get("model_type", "")) if model_config else ""
    if model_type.startswith("deepseek_v4"):
        return True

    if (
        getattr(tokenizer, "tool_call_start", None) == _DEEPSEEK_V4_TOOL_CALL_START
        and getattr(tokenizer, "tool_call_end", None) == _DEEPSEEK_V4_TOOL_CALL_END
    ):
        return True

    return "deepseek-v4" in model_name.lower() or "deepseek_v4" in model_name.lower()


def _serialize_minimax_tool_arguments(arguments: Any) -> str:
    if isinstance(arguments, str):
        return arguments or "{}"
    if arguments is None:
        return "{}"
    try:
        return json.dumps(arguments, ensure_ascii=False, separators=(",", ":"))
    except TypeError:
        return str(arguments)


def _is_minimax_m3_model(
    model_name: str,
    model_config: dict[str, Any] | None = None,
) -> bool:
    model_type = model_config.get("model_type") if model_config else None
    if model_type in _MINIMAX_M3_MODEL_TYPES:
        return True
    lowered = model_name.lower()
    return "minimax" in lowered and "m3" in lowered


class _MiniMaxM3ProtocolNormalizer:
    """Normalize MiniMax M3 protocol markers to oMLX-visible markers."""

    _REPLACEMENTS = (
        (_MINIMAX_THINK_START, "<think>"),
        (_MINIMAX_THINK_END, "</think>"),
        *tuple((token, "") for token in _MINIMAX_SPECIAL_TOKENS),
    )
    _MARKERS = tuple(marker for marker, _ in _REPLACEMENTS)

    def __init__(self) -> None:
        self._buffer = ""

    @classmethod
    def _replace_markers(cls, text: str) -> str:
        for marker, replacement in cls._REPLACEMENTS:
            text = text.replace(marker, replacement)
        return text

    @classmethod
    def _partial_suffix_len(cls, text: str) -> int:
        max_len = min(len(text), max(len(marker) for marker in cls._MARKERS) - 1)
        for size in range(max_len, 0, -1):
            suffix = text[-size:]
            if any(marker.startswith(suffix) for marker in cls._MARKERS):
                return size
        return 0

    def feed(self, text: str) -> str:
        if not text:
            return ""

        self._buffer += text
        keep = self._partial_suffix_len(self._buffer)
        if keep:
            ready = self._buffer[:-keep]
            self._buffer = self._buffer[-keep:]
        else:
            ready = self._buffer
            self._buffer = ""
        return self._replace_markers(ready)

    def finish(self) -> str:
        text = self._replace_markers(self._buffer)
        self._buffer = ""
        return text


def _token_id_for_text(tokenizer: Any, text: str) -> int | None:
    try:
        token_id = tokenizer.convert_tokens_to_ids(text)
    except (AttributeError, KeyError, TypeError, ValueError):
        token_id = None
    if token_id is not None and token_id != getattr(tokenizer, "unk_token_id", None):
        try:
            return int(token_id)
        except (TypeError, ValueError):
            pass

    try:
        token_ids = tokenizer.encode(text, add_special_tokens=False)
    except TypeError:
        try:
            token_ids = tokenizer.encode(text)
        except Exception:
            return None
    except Exception:
        return None

    if len(token_ids) == 1:
        try:
            return int(token_ids[0])
        except (TypeError, ValueError):
            return None
    return None


class DeepSeekV4OutputParserSession:
    """Parser session for DeepSeek V4 DSML tool-call output.

    A completed DSML tool-call block ends the assistant turn. Without a
    parser-owned stop, batched decode keeps the row alive after
    ``</｜DSML｜tool_calls>`` and the model may emit additional or malformed
    DSML fragments as visible assistant text.
    """

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        self._tokenizer = tokenizer
        self._raw_text = ""
        self._stopped = False
        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()

        try:
            from ..api.tool_calling import ToolCallStreamFilter

            self._stream_filter = ToolCallStreamFilter(tokenizer)
            self._visible_filter = ToolCallStreamFilter(tokenizer)
        except Exception as e:  # noqa: BLE001
            logger.debug("DeepSeek V4 stream filter unavailable: %s", e)
            self._stream_filter = None
            self._visible_filter = None

    def _decode_token(self, token_id: int) -> str:
        if self._detokenizer is not None:
            self._detokenizer.add_token(token_id)
            return self._detokenizer.last_segment
        try:
            return self._tokenizer.decode([token_id], skip_special_tokens=False)
        except TypeError:
            return self._tokenizer.decode([token_id])

    def _filtered_text(self, text: str, tool_filter: Any) -> str:
        if not text:
            return ""
        if tool_filter is not None:
            return tool_filter.feed(text)
        return text

    def _finish_filtered_text(self, tool_filter: Any) -> str:
        if tool_filter is None:
            return ""
        return tool_filter.finish()

    def _trim_at_first_tool_block_end(self, text: str) -> tuple[str, bool]:
        start_idx = text.find(_DEEPSEEK_V4_TOOL_CALL_START)
        if start_idx < 0:
            return text, False
        end_idx = text.find(_DEEPSEEK_V4_TOOL_CALL_END, start_idx)
        if end_idx < 0:
            return text, False
        cutoff = end_idx + len(_DEEPSEEK_V4_TOOL_CALL_END)
        return text[:cutoff], True

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        if self._stopped:
            return OutputParserTokenResult(is_stop=True, record_token=False)

        decoded_text = self._decode_token(token_id)
        combined = self._raw_text + decoded_text
        trimmed, is_stop = self._trim_at_first_tool_block_end(combined)

        feed_text = trimmed[len(self._raw_text) :]
        self._raw_text = trimmed
        self._stopped = is_stop

        return OutputParserTokenResult(
            stream_text=self._filtered_text(feed_text, self._stream_filter),
            visible_text=self._filtered_text(feed_text, self._visible_filter),
            is_stop=is_stop,
            record_token=True,
        )

    def finalize(self) -> OutputParserFinalizeResult:
        stream_text = ""
        visible_text = ""
        if self._detokenizer is not None and not self._stopped:
            self._detokenizer.finalize()
            final_text = self._detokenizer.last_segment
            if final_text:
                prev_len = len(self._raw_text)
                combined = self._raw_text + final_text
                self._raw_text, self._stopped = self._trim_at_first_tool_block_end(
                    combined
                )
                final_text = self._raw_text[prev_len:]
                stream_text += self._filtered_text(final_text, self._stream_filter)
                visible_text += self._filtered_text(final_text, self._visible_filter)

        stream_text += self._finish_filtered_text(self._stream_filter)
        visible_text += self._finish_filtered_text(self._visible_filter)

        tool_calls: list[dict[str, str]] = []
        try:
            from ..api.tool_calling import parse_tool_calls

            _, parsed_calls = parse_tool_calls(self._raw_text, self._tokenizer)
            for call in parsed_calls or []:
                tool_calls.append(
                    {
                        "id": getattr(call, "id", ""),
                        "name": call.function.name,
                        "arguments": call.function.arguments,
                    }
                )
        except Exception as e:  # noqa: BLE001
            logger.debug("DeepSeek V4 tool-call parse failed: %s", e)

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
            tool_calls=tool_calls,
            finish_reason="tool_calls" if tool_calls else None,
        )


class MiniMaxM3OutputParserSession:
    """Parser session for MiniMax M3 XML-style tool calls."""

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        self._tokenizer = tokenizer
        self._raw_text = ""
        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()

        try:
            from ..api.tool_calling import ToolCallStreamFilter

            self._stream_filter = ToolCallStreamFilter(tokenizer)
            self._visible_filter = ToolCallStreamFilter(tokenizer)
        except Exception as e:  # noqa: BLE001
            logger.debug("MiniMax M3 stream filter unavailable: %s", e)
            self._stream_filter = None
            self._visible_filter = None
        self._stream_normalizer = _MiniMaxM3ProtocolNormalizer()
        self._visible_normalizer = _MiniMaxM3ProtocolNormalizer()

    def _decode_token(self, token_id: int) -> str:
        if self._detokenizer is not None:
            self._detokenizer.add_token(token_id)
            return self._detokenizer.last_segment
        try:
            return self._tokenizer.decode([token_id], skip_special_tokens=False)
        except TypeError:
            return self._tokenizer.decode([token_id])

    def _filtered_text(
        self,
        text: str,
        tool_filter: Any,
        normalizer: _MiniMaxM3ProtocolNormalizer,
    ) -> str:
        if not text:
            return ""
        if tool_filter is not None:
            text = tool_filter.feed(text)
        return normalizer.feed(text)

    def _finish_filtered_text(
        self,
        tool_filter: Any,
        normalizer: _MiniMaxM3ProtocolNormalizer,
    ) -> str:
        text = ""
        if tool_filter is not None:
            text += normalizer.feed(tool_filter.finish())
        text += normalizer.finish()
        return text

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        decoded_text = self._decode_token(token_id)
        self._raw_text += decoded_text
        is_stop = decoded_text == _MINIMAX_EOS_TOKEN
        return OutputParserTokenResult(
            stream_text=self._filtered_text(
                decoded_text,
                self._stream_filter,
                self._stream_normalizer,
            ),
            visible_text=self._filtered_text(
                decoded_text,
                self._visible_filter,
                self._visible_normalizer,
            ),
            is_stop=is_stop,
            record_token=not is_stop,
        )

    def finalize(self) -> OutputParserFinalizeResult:
        stream_text = ""
        visible_text = ""
        if self._detokenizer is not None:
            self._detokenizer.finalize()
            final_text = self._detokenizer.last_segment
            if final_text:
                self._raw_text += final_text
                stream_text += self._filtered_text(
                    final_text,
                    self._stream_filter,
                    self._stream_normalizer,
                )
                visible_text += self._filtered_text(
                    final_text,
                    self._visible_filter,
                    self._visible_normalizer,
                )

        stream_text += self._finish_filtered_text(
            self._stream_filter,
            self._stream_normalizer,
        )
        visible_text += self._finish_filtered_text(
            self._visible_filter,
            self._visible_normalizer,
        )

        tool_calls: list[dict[str, str]] = []
        if _MINIMAX_TOOL_CALL_START in self._raw_text:
            try:
                from ..patches.mlx_vlm_minimax_m3_compat import (
                    apply_mlx_vlm_minimax_m3_compat_patch,
                )

                apply_mlx_vlm_minimax_m3_compat_patch()

                from mlx_vlm.tool_parsers.minimax_m3 import parse_tool_call

                parsed = parse_tool_call(self._raw_text)
                parsed_calls = parsed if isinstance(parsed, list) else [parsed]
                tool_calls = [
                    {
                        "name": str(call.get("name", "")),
                        "arguments": _serialize_minimax_tool_arguments(
                            call.get("arguments")
                        ),
                    }
                    for call in parsed_calls
                    if isinstance(call, dict) and call.get("name")
                ]
            except Exception as e:  # noqa: BLE001
                logger.debug("MiniMax M3 tool-call parse failed: %s", e)

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
            tool_calls=tool_calls,
            finish_reason="tool_calls" if tool_calls else None,
        )


_INKLING_MODEL_TYPES = {"inkling", "inkling_mm_model"}
_INKLING_MESSAGE_MODEL = "<|message_model|>"
_INKLING_CONTENT_THINKING = "<|content_thinking|>"
_INKLING_CONTENT_TEXT = "<|content_text|>"
_INKLING_CONTENT_XML = "<|content_xml|>"
_INKLING_CONTENT_TOOL_JSON = "<|content_invoke_tool_json|>"
_INKLING_END_MESSAGE = "<|end_message|>"
_INKLING_END_SAMPLING = "<|content_model_end_sampling|>"
_INKLING_MARKERS = (
    _INKLING_MESSAGE_MODEL,
    _INKLING_CONTENT_THINKING,
    _INKLING_CONTENT_TEXT,
    _INKLING_CONTENT_XML,
    _INKLING_CONTENT_TOOL_JSON,
    _INKLING_END_MESSAGE,
    _INKLING_END_SAMPLING,
)


def _is_inkling_model(
    model_name: str,
    model_config: dict[str, Any] | None = None,
) -> bool:
    model_type = model_config.get("model_type") if model_config else None
    if model_type in _INKLING_MODEL_TYPES:
        return True
    return "inkling" in model_name.lower()


class _InklingChannelSplitter:
    """Streaming splitter for inkling's channel protocol.

    The assistant turn is a sequence of blocks::

        [<|message_model|>][HEAD]<|content_*|>BODY<|end_message|> ... \
<|content_model_end_sampling|>

    ``HEAD`` only occurs for tool calls (the function name before
    ``<|content_invoke_tool_json|>``). Thinking bodies surface on the
    stream inside oMLX's ``<think>``/``</think>`` markers, text bodies on
    stream+visible, tool JSON is suppressed (parsed at finalize from the
    raw text).
    """

    def __init__(self) -> None:
        self._buffer = ""
        self._channel: str | None = None
        self._head = ""
        self._think_open = False
        self.stopped = False

    def _partial_suffix_len(self, text: str) -> int:
        max_len = min(len(text), max(len(m) for m in _INKLING_MARKERS) - 1)
        for size in range(max_len, 0, -1):
            suffix = text[-size:]
            if any(m.startswith(suffix) for m in _INKLING_MARKERS):
                return size
        return 0

    def _emit_body(self, text: str) -> tuple[str, str]:
        if not text:
            return "", ""
        if self._channel in ("text", "xml"):
            return text, text
        if self._channel == "thinking":
            # Thinking flows to BOTH channels wrapped in <think> markers
            # (minimax pattern): the scheduler accumulates only
            # visible_text into request.output_text, and the API layer
            # extracts reasoning_content from the <think> block there.
            return text, text
        if self._channel == "tool":
            return "", ""
        # Block head: hold until the next marker classifies it.
        self._head += text
        return "", ""

    def _flush_head_as_text(self) -> tuple[str, str]:
        head, self._head = self._head, ""
        if not head:
            return "", ""
        return head, head

    def _handle_marker(self, marker: str) -> tuple[str, str]:
        stream = visible = ""
        if marker == _INKLING_CONTENT_THINKING:
            s, v = self._flush_head_as_text()
            stream += s
            visible += v
            if not self._think_open:
                stream += "<think>"
                visible += "<think>"
                self._think_open = True
            self._channel = "thinking"
        elif marker in (_INKLING_CONTENT_TEXT, _INKLING_CONTENT_XML):
            s, v = self._flush_head_as_text()
            stream += s
            visible += v
            if self._think_open:
                # A text block after an unterminated thinking block still
                # closes the visible thinking span.
                stream += "</think>"
                visible += "</think>"
                self._think_open = False
            self._channel = "text" if marker == _INKLING_CONTENT_TEXT else "xml"
        elif marker == _INKLING_CONTENT_TOOL_JSON:
            # Head was the tool name; the JSON payload is parsed at
            # finalize from the raw text.
            self._head = ""
            self._channel = "tool"
        elif marker == _INKLING_END_MESSAGE:
            if self._channel == "thinking" and self._think_open:
                stream += "</think>"
                visible += "</think>"
                self._think_open = False
            s, v = self._flush_head_as_text()
            stream += s
            visible += v
            self._channel = None
        elif marker == _INKLING_MESSAGE_MODEL:
            s, v = self._flush_head_as_text()
            stream += s
            visible += v
            self._channel = None
        elif marker == _INKLING_END_SAMPLING:
            self.stopped = True
            self._channel = None
        return stream, visible

    def feed(self, text: str) -> tuple[str, str]:
        if not text:
            return "", ""
        self._buffer += text
        stream = visible = ""
        while True:
            first_idx = -1
            first_marker = None
            for marker in _INKLING_MARKERS:
                idx = self._buffer.find(marker)
                if idx >= 0 and (first_idx < 0 or idx < first_idx):
                    first_idx = idx
                    first_marker = marker
            if first_marker is None:
                break
            s, v = self._emit_body(self._buffer[:first_idx])
            stream += s
            visible += v
            m_s, m_v = self._handle_marker(first_marker)
            stream += m_s
            visible += m_v
            self._buffer = self._buffer[first_idx + len(first_marker) :]

        keep = self._partial_suffix_len(self._buffer)
        ready = self._buffer[: len(self._buffer) - keep]
        self._buffer = self._buffer[len(self._buffer) - keep :]
        s, v = self._emit_body(ready)
        return stream + s, visible + v

    def finish(self) -> tuple[str, str]:
        stream = visible = ""
        s, v = self._emit_body(self._buffer)
        stream += s
        visible += v
        self._buffer = ""
        s, v = self._flush_head_as_text()
        stream += s
        visible += v
        if self._think_open:
            stream += "</think>"
            visible += "</think>"
            self._think_open = False
        return stream, visible


class InklingOutputParserSession:
    """Parser session for inkling channel output (thinking / text / tool)."""

    _TOOL_RE = None  # compiled lazily

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        import re

        self._tokenizer = tokenizer
        self._raw_text = ""
        self._splitter = _InklingChannelSplitter()
        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()
        if InklingOutputParserSession._TOOL_RE is None:
            InklingOutputParserSession._TOOL_RE = re.compile(
                re.escape(_INKLING_CONTENT_TOOL_JSON)
                + r"(.*?)(?:"
                + re.escape(_INKLING_END_MESSAGE)
                + r"|"
                + re.escape(_INKLING_END_SAMPLING)
                + r"|\Z)",
                re.S,
            )

    def _decode_token(self, token_id: int) -> str:
        if self._detokenizer is not None:
            self._detokenizer.add_token(token_id)
            return self._detokenizer.last_segment
        try:
            return self._tokenizer.decode([token_id], skip_special_tokens=False)
        except TypeError:
            return self._tokenizer.decode([token_id])

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        if self._splitter.stopped:
            return OutputParserTokenResult(is_stop=True, record_token=False)
        decoded_text = self._decode_token(token_id)
        self._raw_text += decoded_text
        stream_text, visible_text = self._splitter.feed(decoded_text)
        is_stop = self._splitter.stopped
        return OutputParserTokenResult(
            stream_text=stream_text,
            visible_text=visible_text,
            is_stop=is_stop,
            record_token=not is_stop,
        )

    def finalize(self) -> OutputParserFinalizeResult:
        stream_text = ""
        visible_text = ""
        if self._detokenizer is not None and not self._splitter.stopped:
            self._detokenizer.finalize()
            final_text = self._detokenizer.last_segment
            if final_text:
                self._raw_text += final_text
                s, v = self._splitter.feed(final_text)
                stream_text += s
                visible_text += v
        s, v = self._splitter.finish()
        stream_text += s
        visible_text += v

        tool_calls: list[dict[str, str]] = []
        for match in InklingOutputParserSession._TOOL_RE.finditer(self._raw_text):
            payload = match.group(1).strip()
            if not payload:
                continue
            try:
                parsed = json.loads(payload)
            except (json.JSONDecodeError, ValueError):
                logger.debug("Inkling tool-call payload not valid JSON")
                continue
            if not isinstance(parsed, dict) or not parsed.get("name"):
                continue
            args = parsed.get("args", {})
            tool_calls.append(
                {
                    "name": str(parsed["name"]),
                    "arguments": json.dumps(
                        args if isinstance(args, dict) else {},
                        ensure_ascii=False,
                        separators=(",", ":"),
                    ),
                }
            )

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
            tool_calls=tool_calls,
            finish_reason="tool_calls" if tool_calls else None,
        )


def _create_cohere2_moe_filter():
    try:
        from cohere_melody import PyFilter, PyFilterOptions
    except ImportError:
        return None

    return PyFilter(PyFilterOptions().cmd4().stream_tool_actions())


def _reserialize_cohere_tool_arguments(args: str) -> str:
    if not args:
        return "{}"
    try:
        return json.dumps(
            json.loads(args, strict=False),
            ensure_ascii=False,
            separators=(",", ":"),
        )
    except (json.JSONDecodeError, ValueError):
        return args or "{}"


class Cohere2MoeOutputParserSession:
    """Parser session for Cohere2 MoE / Command-style Melody output."""

    def __init__(self, tokenizer: Any, model_path: str | None = None):
        self._tokenizer = tokenizer
        self._melody = _create_cohere2_moe_filter()
        if self._melody is None:
            raise RuntimeError("cohere_melody is not installed")

        self._detokenizer = create_streaming_detokenizer(tokenizer, model_path)
        if self._detokenizer is not None:
            self._detokenizer.reset()

        self._thinking_started = False
        self._thinking_closed = False
        self._tool_calls: dict[int, dict[str, str]] = {}

    def _decode_token(self, token_id: int) -> str:
        if self._detokenizer is not None:
            self._detokenizer.add_token(token_id)
            return self._detokenizer.last_segment
        try:
            return self._tokenizer.decode([token_id], skip_special_tokens=False)
        except TypeError:
            return self._tokenizer.decode([token_id])

    def _accumulate_tool_calls(self, tool_calls: list[Any]) -> None:
        for tool_call in tool_calls:
            index = int(getattr(tool_call, "index", 0) or 0)
            current = self._tool_calls.setdefault(
                index,
                {"id": "", "name": "", "arguments": ""},
            )
            current["id"] += getattr(tool_call, "id", "") or ""
            current["name"] += getattr(tool_call, "name", "") or ""
            current["arguments"] += getattr(tool_call, "arguments", "") or ""

    def _apply_melody_result(self, result: Any) -> tuple[str, str]:
        stream_text = ""
        visible_text = ""

        reasoning = getattr(result, "reasoning", None)
        if reasoning:
            if not self._thinking_started:
                self._thinking_started = True
                stream_text += "<think>\n"
                visible_text += "<think>\n"
            stream_text += reasoning
            visible_text += reasoning

        content = getattr(result, "content", None)
        if content:
            if self._thinking_started and not self._thinking_closed:
                self._thinking_closed = True
                stream_text += "</think>\n"
                visible_text += "</think>\n"
            stream_text += content
            visible_text += content

        self._accumulate_tool_calls(getattr(result, "tool_calls", []) or [])
        return stream_text, visible_text

    def process_token(self, token_id: int) -> OutputParserTokenResult:
        decoded_text = self._decode_token(token_id)
        if not decoded_text:
            return OutputParserTokenResult(record_token=True)

        result = self._melody.write_decoded(decoded_text)
        stream_text, visible_text = self._apply_melody_result(result)
        return OutputParserTokenResult(
            stream_text=stream_text,
            visible_text=visible_text,
            record_token=True,
        )

    def finalize(self) -> OutputParserFinalizeResult:
        stream_text = ""
        visible_text = ""

        if self._detokenizer is not None:
            self._detokenizer.finalize()
            final_text = self._detokenizer.last_segment
            if final_text:
                result = self._melody.write_decoded(final_text)
                s_text, v_text = self._apply_melody_result(result)
                stream_text += s_text
                visible_text += v_text

        result = self._melody.flush_partials()
        s_text, v_text = self._apply_melody_result(result)
        stream_text += s_text
        visible_text += v_text

        if self._thinking_started and not self._thinking_closed:
            self._thinking_closed = True
            stream_text += "</think>\n"
            visible_text += "</think>\n"

        tool_calls = [
            {
                "id": value["id"],
                "name": value["name"],
                "arguments": _reserialize_cohere_tool_arguments(value["arguments"]),
            }
            for _, value in sorted(self._tool_calls.items())
            if value["name"]
        ]

        return OutputParserFinalizeResult(
            stream_text=stream_text,
            visible_text=visible_text,
            tool_calls=tool_calls,
            finish_reason="tool_calls" if tool_calls else None,
        )


def detect_output_parser(
    model_name: str,
    tokenizer: Any,
    model_config: dict[str, Any] | None = None,
    model_path: str | None = None,
) -> OutputParserFactory | None:
    """Detect a protocol-specific output parser for the model, if needed.

    ``model_name`` drives detection (string matching) and may be a display
    id rather than a directory since #2178. Pass ``model_path`` when the
    filesystem path is available so parser sessions can locate
    tokenizer.json for their streaming detokenizers.
    """
    session_model_path = model_path or model_name

    if is_harmony_model(model_name, model_config):
        temp_parser = HarmonyStreamingParser(tokenizer)
        return OutputParserFactory(
            kind="harmony",
            create_session=lambda session_tokenizer: HarmonyOutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=temp_parser.get_stop_token_ids(),
            thinking_end_text="<|end|>",
            thinking_end_trailing_text="<|start|>assistant<|channel|>final<|message|>",
        )

    if is_gemma4_model(model_name, model_config):
        from .gemma4 import (
            _CLOSE_MARKER,
            _OPEN_MARKER_BARE,
            _TOOL_RESPONSE_CLOSE,
            _TOOL_RESPONSE_OPEN,
            _TURN_END_MARKER,
            Gemma4OutputParserSession,
        )

        return OutputParserFactory(
            kind="gemma4",
            create_session=lambda session_tokenizer: Gemma4OutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=set(),
            thinking_end_text="<channel|>",
            protocol_marker_texts=(
                _OPEN_MARKER_BARE,
                _CLOSE_MARKER,
                _TURN_END_MARKER,
                _TOOL_RESPONSE_OPEN,
                _TOOL_RESPONSE_CLOSE,
            ),
        )

    if _is_deepseek_v4_model(model_name, tokenizer, model_config):
        return OutputParserFactory(
            kind="deepseek_v4",
            create_session=lambda session_tokenizer: DeepSeekV4OutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=set(),
            protocol_marker_texts=(
                _DEEPSEEK_V4_TOOL_CALL_START,
                _DEEPSEEK_V4_TOOL_CALL_END,
            ),
        )

    if _is_cohere2_moe_model(model_name, model_config):
        if _create_cohere2_moe_filter() is None:
            logger.warning(
                "cohere_melody is not installed; Cohere2 MoE output parser "
                "is disabled for %s",
                model_name,
            )
            return None

        return OutputParserFactory(
            kind="cohere2_moe",
            create_session=lambda session_tokenizer: Cohere2MoeOutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=set(),
            thinking_end_text="</think>",
        )

    if _is_inkling_model(model_name, model_config):
        inkling_stop_ids = set()
        end_sampling_id = _token_id_for_text(tokenizer, _INKLING_END_SAMPLING)
        if end_sampling_id is not None:
            inkling_stop_ids.add(end_sampling_id)

        return OutputParserFactory(
            kind="inkling",
            create_session=lambda session_tokenizer: InklingOutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=inkling_stop_ids,
            thinking_start_text=_INKLING_CONTENT_THINKING,
            thinking_start_output_text="<think>\n",
            thinking_end_text=_INKLING_END_MESSAGE,
            thinking_end_trailing_text=(
                _INKLING_MESSAGE_MODEL + _INKLING_CONTENT_TEXT
            ),
            protocol_marker_texts=_INKLING_MARKERS,
        )

    if _is_minimax_m3_model(model_name, model_config):
        minimax_stop_ids = set()
        eos_id = _token_id_for_text(tokenizer, _MINIMAX_EOS_TOKEN)
        if eos_id is not None:
            minimax_stop_ids.add(eos_id)

        return OutputParserFactory(
            kind="minimax_m3",
            create_session=lambda session_tokenizer: MiniMaxM3OutputParserSession(
                session_tokenizer,
                model_path=session_model_path,
            ),
            stop_token_ids=minimax_stop_ids,
            thinking_start_text=_MINIMAX_THINK_START,
            thinking_start_output_text="<think>\n",
            thinking_end_text=_MINIMAX_THINK_END,
            protocol_marker_texts=(
                _MINIMAX_THINK_START,
                _MINIMAX_THINK_END,
                _MINIMAX_TOOL_CALL_START,
                _MINIMAX_TOOL_CALL_END,
            ),
        )

    return None


def detect_message_extractor(
    model_name: str,
    model_config: dict[str, Any] | None = None,
) -> Callable:
    """Return the appropriate message extractor function for the model.

    The returned callable has the signature::

        extractor(messages, max_tool_result_tokens=None, tokenizer=None) -> list[dict]

    This mirrors how ``detect_output_parser`` decouples model-specific
    knowledge from the server layer — the engine stores the extractor at
    load time and the server just calls ``engine.message_extractor(...)``.
    """
    if is_harmony_model(model_name, model_config):
        from ..api.utils import extract_harmony_messages

        return extract_harmony_messages

    if is_gemma4_model(model_name, model_config):
        from .gemma4 import extract_gemma4_messages

        return extract_gemma4_messages

    # Default: caller decides between extract_text_content and
    # extract_multimodal_content based on engine type (VLM vs text).
    return None


# ===== FILE: omlx/admin/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""Admin panel for oMLX server configuration."""

from .auth import create_session_token, require_admin, verify_session
from .routes import router as admin_router, set_admin_getters, set_hf_downloader

__all__ = [
    "admin_router",
    "create_session_token",
    "require_admin",
    "set_admin_getters",
    "set_hf_downloader",
    "verify_session",
]


# ===== FILE: omlx/admin/accuracy_benchmark.py =====

# SPDX-License-Identifier: Apache-2.0
"""Accuracy benchmark execution logic for oMLX admin panel.

Orchestrates MMLU, HellaSwag, TruthfulQA, GSM8K, and LiveCodeBench
evaluations with real-time progress reporting via SSE events.

Supports server-side queue and persistent result accumulation.
Results survive browser close and persist until explicitly reset.
"""

import asyncio
import logging
import time
import uuid
from collections import Counter
from dataclasses import dataclass, field
from typing import Any, Literal, Optional

from pydantic import BaseModel, field_validator, model_validator

from .external_api import (
    ExternalAPIClient,
    ExternalChatAdapter,
    ExternalEndpointConfig,
)

logger = logging.getLogger(__name__)

# Module-level storage for active benchmark runs
_accuracy_runs: dict[str, "AccuracyBenchmarkRun"] = {}

# Accumulated results — persists until explicit reset
_accumulated_results: list[dict] = []

# Server-side queue
_queue: list["AccuracyBenchmarkRequest"] = []
_queue_running: bool = False
_current_run_id: Optional[str] = None
_current_model: Optional[str] = None
_engine_pool_ref: Any = None
# Chain-ownership token. A "chain" is one start_next_from_queue call plus
# the _continue_queue tail it spawns. Each chain captures the token current
# at its start; cancel_queue and start_next_from_queue bump it. A chain
# whose token is stale (e.g. it was soft-cancelled and only noticed at its
# next checkpoint, after the user already started a new chain) must not pop
# the queue or mutate _queue_running/_current_run_id — otherwise it starts
# a run concurrently with the live chain, whose Phase 1 "unload all models"
# rips the engine out from under the active run.
_chain_id: int = 0

VALID_BENCHMARKS = [
    "mmlu", "mmlu_pro", "kmmlu", "cmmlu", "jmmlu",
    "hellaswag", "truthfulqa", "arc_challenge", "winogrande",
    "gsm8k", "mathqa", "humaneval", "mbpp", "livecodebench",
    "bbq", "safetybench",
]

# Sampling profile for an accuracy run. "deterministic" (default) runs greedy
# (temperature 0) so saved scores stay reproducible; "model_settings" opts in to
# the model's configured sampling (temperature, top_p, …) for a real-world score.
SamplingProfile = Literal["deterministic", "model_settings"]


class AccuracyBenchmarkRequest(BaseModel):
    """Request model for starting an accuracy benchmark."""

    model_id: str
    benchmarks: dict[str, int]  # name -> sample_size (0 = full dataset)
    batch_size: int = 1
    enable_thinking: bool = False
    sampling_profile: SamplingProfile = "deterministic"
    # When set, the benchmark runs against a remote OpenAI-compatible
    # endpoint instead of a local engine and model_id is the remote
    # model name (not validated against the local catalog).
    external: Optional[ExternalEndpointConfig] = None

    @model_validator(mode="after")
    def _force_thinking_off_for_external(self) -> "AccuracyBenchmarkRequest":
        # enable_thinking is a local chat-template kwarg; external requests
        # never send it, so keep the stored flag honest.
        if self.external is not None:
            self.enable_thinking = False
        return self

    @field_validator("batch_size")
    @classmethod
    def validate_batch_size(cls, v: int) -> int:
        if v not in (1, 2, 4, 8, 16, 32):
            raise ValueError("batch_size must be 1, 2, 4, 8, 16, or 32")
        return v

    @field_validator("benchmarks")
    @classmethod
    def validate_benchmarks(cls, v: dict[str, int]) -> dict[str, int]:
        if not v:
            raise ValueError("At least one benchmark is required")
        for name, size in v.items():
            if name not in VALID_BENCHMARKS:
                raise ValueError(
                    f"Invalid benchmark '{name}'. Must be one of {VALID_BENCHMARKS}"
                )
            if size < 0:
                raise ValueError(f"Sample size for '{name}' must be >= 0")
        return v


@dataclass
class AccuracyBenchmarkRun:
    """Tracks the state of a running accuracy benchmark.

    SSE delivery model mirrors `BenchmarkRun`: append-only `events`
    log + `cond` for live notification + `terminal` flag set on the
    final event. See benchmark.py for the rationale.
    """

    bench_id: str
    request: AccuracyBenchmarkRequest
    status: str = "running"  # running, completed, cancelled, error
    events: list[dict] = field(default_factory=list)
    cond: asyncio.Condition = field(default_factory=asyncio.Condition)
    terminal: bool = False
    task: Optional[asyncio.Task] = None
    results: list[dict] = field(default_factory=list)
    error_message: str = ""
    last_progress: Optional[dict] = None  # last progress event for reconnect
    # Finer-grained lifecycle than `status` — surfaces the difference between
    # "still scoring questions" and "cleaning up after the last result was
    # emitted". The serialization gate (_queue_running) stays True across
    # both, but a UI rendering the running row wants to hide it once
    # phase=="unloading" so the user isn't told "still running" when the
    # result card has already appeared on screen. Transitions:
    #   pending → loading → evaluating → unloading → completed
    # (cancelled / error replace the terminal phase on those branches.)
    phase: str = "pending"


# Accuracy stream closes on `done` (run finished) or `error`. Unlike the
# throughput bench there's no separate upload phase to ride out.
_ACCURACY_TERMINAL_TYPES = frozenset({"done", "error"})


# --- Run management ---


def get_run(bench_id: str) -> Optional[AccuracyBenchmarkRun]:
    """Get an accuracy benchmark run by ID."""
    return _accuracy_runs.get(bench_id)


def create_run(request: AccuracyBenchmarkRequest) -> AccuracyBenchmarkRun:
    """Create a new accuracy benchmark run."""
    bench_id = str(uuid.uuid4())[:8]
    run = AccuracyBenchmarkRun(bench_id=bench_id, request=request)
    _accuracy_runs[bench_id] = run
    return run


def cleanup_old_runs() -> None:
    """Remove completed/errored runs to prevent memory leaks."""
    to_remove = []
    for bid, run in _accuracy_runs.items():
        if run.status in ("completed", "cancelled", "error"):
            to_remove.append(bid)
    for bid in to_remove:
        del _accuracy_runs[bid]


# --- Accumulated results ---


def get_accumulated_results() -> list[dict]:
    """Get all accumulated benchmark results."""
    return _accumulated_results


def reset_accumulated_results() -> None:
    """Clear all accumulated results."""
    _accumulated_results.clear()


# --- Queue management ---


def add_to_queue(request: AccuracyBenchmarkRequest) -> None:
    """Add a benchmark request to the queue."""
    _queue.append(request)


def get_queue_status() -> dict:
    """Get current queue status."""
    last_progress = None
    phase = None
    if _current_run_id:
        run = get_run(_current_run_id)
        if run:
            last_progress = run.last_progress
            phase = run.phase
    return {
        "running": _queue_running,
        "current_model": _current_model,
        "current_bench_id": _current_run_id,
        "last_progress": last_progress,
        # Finer-grained than `running`: distinguishes "still scoring" from
        # "cleaning up after the last result emitted". Polling UIs hide
        # the running row once phase becomes "unloading" / "completed" so
        # the result card alone tells the story.
        "phase": phase,
        "queue": [
            {
                "model_id": r.model_id,
                "benchmarks": list(r.benchmarks.keys()),
                "external": r.external is not None,
            }
            for r in _queue
        ],
    }


def remove_from_queue(idx: int) -> bool:
    """Remove an item from the queue by index."""
    if 0 <= idx < len(_queue):
        _queue.pop(idx)
        return True
    return False


def start_next_from_queue(engine_pool: Any) -> Optional[str]:
    """Pop next item from queue, create run, start background task.

    Returns bench_id if a run was started, None if already running or queue empty.
    This is synchronous so the caller gets the bench_id immediately.
    """
    global _queue_running, _current_run_id, _current_model, _engine_pool_ref
    global _chain_id

    _engine_pool_ref = engine_pool

    if _queue_running:
        return None

    if not _queue:
        return None

    request = _queue.pop(0)
    _queue_running = True
    _current_model = request.model_id
    # This chain takes ownership of the queue; any earlier chain still
    # draining a soft-cancelled run bails at its next _continue_queue call.
    _chain_id += 1
    my_chain = _chain_id

    cleanup_old_runs()
    run = create_run(request)
    _current_run_id = run.bench_id

    logger.info(
        f"Queue: starting {request.model_id} "
        f"benchmarks={list(request.benchmarks.keys())}"
    )

    async def _run_and_continue():
        try:
            await run_accuracy_benchmark(run, engine_pool)
        except Exception as e:
            logger.error(f"Queue: error running {request.model_id}: {e}")
        # Auto-continue with next in queue
        await _continue_queue(engine_pool, my_chain)

    run.task = asyncio.create_task(_run_and_continue())
    return run.bench_id


async def _continue_queue(engine_pool: Any, chain_id: int) -> None:
    """Continue processing the queue after a run completes.

    `chain_id` is the ownership token captured when this chain started.
    A stale chain (orphaned by cancel_queue, with a new chain started by
    the user since) returns without popping the queue or touching the
    gate, so it cannot start a run concurrently with the live chain.
    """
    global _queue_running, _current_run_id, _current_model

    if chain_id != _chain_id:
        return

    if not _queue:
        _queue_running = False
        _current_run_id = None
        _current_model = None
        return

    request = _queue.pop(0)
    _current_model = request.model_id

    cleanup_old_runs()
    run = create_run(request)
    _current_run_id = run.bench_id
    # Queue-continued runs execute inside this chain's own task; record it
    # so cancel_queue can hard-cancel them instead of waiting for the next
    # on_progress checkpoint (up to a full generation batch away).
    run.task = asyncio.current_task()

    logger.info(
        f"Queue: continuing with {request.model_id} "
        f"benchmarks={list(request.benchmarks.keys())}"
    )

    try:
        await run_accuracy_benchmark(run, engine_pool)
    except Exception as e:
        logger.error(f"Queue: error running {request.model_id}: {e}")

    await _continue_queue(engine_pool, chain_id)


async def cancel_queue() -> None:
    """Cancel the current run and clear the queue."""
    global _queue_running, _current_run_id, _current_model, _chain_id

    _queue.clear()
    # Orphan the live chain before releasing the gate: if the cancelled run
    # only notices at its next checkpoint, its trailing _continue_queue must
    # not race whatever chain the user starts after this cancel.
    _chain_id += 1

    if _current_run_id:
        run = get_run(_current_run_id)
        if run and run.status == "running":
            run.status = "cancelled"
            if run.task and not run.task.done():
                run.task.cancel()

    _queue_running = False
    _current_run_id = None
    _current_model = None


# --- SSE ---


async def _send_event(run: AccuracyBenchmarkRun, event: dict) -> None:
    """Append an event to the run's log and wake subscribers.

    Updates `last_progress` (used by the REST `queue/status` endpoint
    for reconnect hints) and sets `run.terminal` on the final event.
    """
    if event.get("type") == "progress":
        run.last_progress = event
    async with run.cond:
        run.events.append(event)
        if event.get("type") in _ACCURACY_TERMINAL_TYPES:
            run.terminal = True
        run.cond.notify_all()


# --- Benchmark execution ---


async def run_accuracy_benchmark(
    run: AccuracyBenchmarkRun, engine_pool: Any
) -> None:
    """Execute accuracy benchmark run.

    Phases:
    1. Unload all models
    2. Load target model
    3. For each selected benchmark: load data, evaluate, report
    4. Unload model
    5. Send done event
    """
    from ..eval import BENCHMARKS

    request = run.request

    # Suppress TTL auto-unload during benchmark (local engines only)
    if request.external is None:
        engine_pool._suppress_ttl = True
    start_time = time.time()
    client: Optional[ExternalAPIClient] = None

    try:
        run.phase = "loading"
        if request.external is not None:
            # External endpoint: no local model lifecycle. The adapter owns
            # the sampling-profile mapping, so sampling_kwargs stays empty
            # (enable_thinking is already forced off by request validation).
            await _send_event(run, {
                "type": "progress",
                "phase": "connect",
                "model_id": request.model_id,
                "benchmark": "",
                "message": f"Connecting to {request.external.base_url}...",
                "current": 0,
                "total": len(request.benchmarks),
            })
            client = ExternalAPIClient(request.external)
            engine = ExternalChatAdapter(client, request.sampling_profile)
            # Fail fast on auth/URL/model errors so a wrong API key cannot
            # silently produce a 0% score.
            await engine.preflight()
            sampling_kwargs = {}
        else:
            # Phase 1: Unload all models
            loaded_ids = engine_pool.get_loaded_model_ids()
            if loaded_ids:
                await _send_event(run, {
                    "type": "progress",
                    "phase": "unload",
                    "model_id": request.model_id,
                    "benchmark": "",
                    "message": f"Unloading {len(loaded_ids)} model(s)...",
                    "current": 0,
                    "total": len(request.benchmarks),
                })
                for model_id in loaded_ids:
                    try:
                        await engine_pool._unload_engine(model_id)
                    except Exception as e:
                        logger.warning(f"Failed to unload {model_id}: {e}")

            # Phase 2: Load target model
            await _send_event(run, {
                "type": "progress",
                "phase": "load",
                "model_id": request.model_id,
                "benchmark": "",
                "message": f"Loading {request.model_id}...",
                "current": 0,
                "total": len(request.benchmarks),
            })

            # Force LM engine for accuracy benchmarks — text-only tasks
            # don't need VLM and the VLM adapter can produce empty responses.
            engine = await engine_pool.get_engine(request.model_id, force_lm=True)

            # Load model sampling settings. Under the default "deterministic"
            # profile sampling params are not read — the benchmark runs greedy
            # (temperature 0) so saved scores stay reproducible. Only the
            # explicit "model_settings" opt-in honors the model's configured
            # sampling. chat_template_kwargs is prompt construction, not
            # sampling, so it is forwarded in both profiles.
            sampling_kwargs = {}
            if engine_pool._settings_manager is not None:
                ms = engine_pool._settings_manager.get_settings(request.model_id)
                if ms.chat_template_kwargs:
                    sampling_kwargs["chat_template_kwargs"] = ms.chat_template_kwargs
                if request.sampling_profile == "model_settings":
                    if ms.temperature is not None:
                        sampling_kwargs["temperature"] = ms.temperature
                    if ms.top_p is not None:
                        sampling_kwargs["top_p"] = ms.top_p
                    if ms.top_k is not None:
                        sampling_kwargs["top_k"] = ms.top_k
                    if ms.min_p is not None:
                        sampling_kwargs["min_p"] = ms.min_p
                    if ms.repetition_penalty is not None:
                        sampling_kwargs["repetition_penalty"] = ms.repetition_penalty
                    if ms.presence_penalty is not None:
                        sampling_kwargs["presence_penalty"] = ms.presence_penalty

        # Phase 3: Run each benchmark
        run.phase = "evaluating"
        completed = 0
        for bench_name, sample_size in request.benchmarks.items():
            if run.status == "cancelled":
                break

            bench_cls = BENCHMARKS.get(bench_name)
            if bench_cls is None:
                logger.warning(f"Unknown benchmark: {bench_name}")
                continue

            evaluator = bench_cls()

            # Load dataset
            await _send_event(run, {
                "type": "progress",
                "phase": "download",
                "model_id": request.model_id,
                "benchmark": bench_name,
                "message": f"Loading {bench_name} dataset...",
                "current": completed,
                "total": len(request.benchmarks),
            })

            try:
                items = await evaluator.load_dataset(sample_size=sample_size)
            except Exception as e:
                logger.error(f"Failed to load {bench_name} dataset: {e}")
                await _send_event(run, {
                    "type": "error",
                    "message": f"Failed to load {bench_name} dataset: {e}",
                })
                run.status = "error"
                run.error_message = str(e)
                return

            # Run evaluation with progress
            total_items = len(items)

            async def on_progress(current: int, total: int) -> None:
                if run.status == "cancelled":
                    raise asyncio.CancelledError()
                await _send_event(run, {
                    "type": "progress",
                    "phase": "eval",
                    "model_id": request.model_id,
                    "benchmark": bench_name,
                    "message": f"Evaluating {bench_name} ({current}/{total})...",
                    "current": completed,
                    "total": len(request.benchmarks),
                    "bench_current": current,
                    "bench_total": total,
                })

            await _send_event(run, {
                "type": "progress",
                "phase": "eval",
                "model_id": request.model_id,
                "benchmark": bench_name,
                "message": f"Evaluating {bench_name} (0/{total_items})...",
                "current": completed,
                "total": len(request.benchmarks),
                "bench_current": 0,
                "bench_total": total_items,
            })

            try:
                result = await evaluator.run(
                    engine, items, on_progress,
                    batch_size=request.batch_size,
                    sampling_kwargs=sampling_kwargs,
                    enable_thinking=request.enable_thinking,
                )
            except asyncio.CancelledError:
                run.status = "cancelled"
                await _send_event(run, {
                    "type": "error",
                    "message": "Benchmark cancelled",
                })
                return
            except Exception as e:
                logger.error(f"Error running {bench_name}: {e}")
                await _send_event(run, {
                    "type": "error",
                    "message": f"Error running {bench_name}: {e}",
                })
                run.status = "error"
                run.error_message = str(e)
                return

            question_results = []
            for qr in result.question_results:
                question_data = {
                    "id": qr.question_id,
                    "correct": qr.correct,
                    "expected": qr.expected,
                    "predicted": qr.predicted,
                    "question": qr.question_text,
                    "raw_response": qr.raw_response,
                    "category": qr.category,
                    "time_s": round(qr.time_seconds, 3),
                }
                if request.external is not None:
                    question_data.update({
                        "status": qr.status,
                        "finish_reason": qr.finish_reason,
                        "reasoning_fields_present": qr.reasoning_fields_present,
                        "reasoning_fields_nonempty": qr.reasoning_fields_nonempty,
                        "prompt_tokens": qr.prompt_tokens,
                        "completion_tokens": qr.completion_tokens,
                        "error_message": qr.error_message,
                    })
                question_results.append(question_data)

            result_data = {
                "model_id": request.model_id,
                "external": request.external is not None,
                "benchmark": result.benchmark_name,
                "accuracy": round(result.accuracy, 4),
                "thinking_used": result.thinking_used,
                "total": result.total_questions,
                "correct": result.correct_count,
                "time_s": round(result.time_seconds, 1),
                "question_results": question_results,
            }
            if request.external is not None:
                status_counts = Counter(
                    qr.status or "invalid_response" for qr in result.question_results
                )
                valid_responses = status_counts["correct"] + status_counts["wrong"]
                total_questions = result.total_questions
                result_data.update({
                    "valid_response_count": valid_responses,
                    "empty_content_count": status_counts["empty_content"],
                    "truncated_count": status_counts["truncated"],
                    "timeout_count": status_counts["timeout"],
                    "http_error_count": status_counts["http_error"],
                    "connection_error_count": status_counts["connection_error"],
                    "invalid_response_count": status_counts["invalid_response"],
                    "parse_error_count": status_counts["parse_error"],
                    "wrong_count": status_counts["wrong"],
                    "valid_response_rate": round(
                        valid_responses / total_questions
                        if total_questions > 0 else 0.0,
                        4,
                    ),
                    "valid_answer_accuracy": round(
                        result.correct_count / valid_responses
                        if valid_responses > 0 else 0.0,
                        4,
                    ),
                    "reliability_warning": any(
                        status_counts[name] > 0
                        for name in (
                            "empty_content",
                            "truncated",
                            "timeout",
                            "http_error",
                            "connection_error",
                            "invalid_response",
                            "parse_error",
                        )
                    ),
                })
            if result.category_scores:
                result_data["category_scores"] = {
                    k: round(v, 4) for k, v in result.category_scores.items()
                }

            # Accumulate persistently
            _accumulated_results.append(result_data)

            run.results.append(result_data)
            completed += 1

            await _send_event(run, {
                "type": "result",
                "data": result_data,
            })

        # Phase 4: Unload model. The result(s) are already emitted by now,
        # so flip phase so polling clients hide the running indicator
        # (the result card has already appeared on screen — telling the
        # user "still running" while we clean up reads as a bug).
        run.phase = "unloading"
        if request.external is None:
            try:
                await engine_pool._unload_engine(request.model_id)
            except Exception:
                pass

        # Phase 5: Done
        total_time = time.time() - start_time
        run.status = "completed"
        run.phase = "completed"

        await _send_event(run, {
            "type": "done",
            "summary": {
                "model_id": request.model_id,
                "total_time": round(total_time, 1),
                "benchmarks_completed": completed,
            },
        })

    except asyncio.CancelledError:
        run.status = "cancelled"
        run.phase = "cancelled"
        await _send_event(run, {
            "type": "error",
            "message": "Benchmark cancelled",
        })
    except Exception as e:
        logger.exception(f"Accuracy benchmark error: {e}")
        run.status = "error"
        run.phase = "error"
        run.error_message = str(e)
        await _send_event(run, {
            "type": "error",
            "message": str(e),
        })
    finally:
        # Re-enable TTL auto-unload
        engine_pool._suppress_ttl = False
        if client is not None:
            await client.aclose()


# ===== FILE: omlx/admin/auth.py =====

# SPDX-License-Identifier: Apache-2.0
"""Authentication utilities for the oMLX admin panel.

This module provides session-based authentication using signed tokens
and API key verification for admin panel access.
"""

import hashlib
import os
import secrets
from typing import Optional

from fastapi import HTTPException, Request
from fastapi.responses import RedirectResponse
from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

# Session configuration
SESSION_COOKIE_NAME = "omlx_admin_session"
SESSION_MAX_AGE = 86400  # 24 hours in seconds
REMEMBER_ME_MAX_AGE = 2592000  # 30 days in seconds

# Secret key for signing session tokens
# Use environment variable if set, otherwise generate a random key
# Note: Random key means sessions won't persist across server restarts
# This is a fallback; init_auth() should be called with a persistent key
SECRET_KEY = os.environ.get("OMLX_SECRET_KEY") or secrets.token_hex(32)

# Initialize the serializer for creating and verifying session tokens
_serializer = URLSafeTimedSerializer(SECRET_KEY)

# Global settings getter (set by init_auth)
_get_global_settings = None


def init_auth(secret_key: str, global_settings_getter=None) -> None:
    """Initialize authentication with a persistent secret key.

    Should be called during server startup with the secret key from settings.
    Environment variable OMLX_SECRET_KEY takes priority if set.

    Args:
        secret_key: The secret key from settings.json for signing tokens.
        global_settings_getter: Optional callable that returns GlobalSettings.
    """
    global _serializer, SECRET_KEY, _get_global_settings
    # Environment variable takes priority over settings
    key = os.environ.get("OMLX_SECRET_KEY") or secret_key
    SECRET_KEY = key
    _serializer = URLSafeTimedSerializer(key)
    if global_settings_getter is not None:
        _get_global_settings = global_settings_getter


def create_session_token(remember: bool = False) -> str:
    """Create a signed session token for admin authentication.

    Args:
        remember: If True, the token payload includes a remember flag
                  for extended session duration (30 days).

    Returns:
        A URL-safe signed token string containing admin session data.

    Example:
        >>> token = create_session_token()
        >>> verify_session_token(token)
        True
    """
    payload = {"admin": True, "remember": remember}
    return _serializer.dumps(payload)


def verify_session_token(token: str, max_age: int = SESSION_MAX_AGE) -> bool:
    """Verify and decode a session token.

    The max_age is determined by the token's remember flag:
    - remember=True: 30 days
    - remember=False (default): 24 hours

    Args:
        token: The signed session token to verify.
        max_age: Maximum age of the token in seconds. Defaults to 24 hours.
                 This is overridden by the token's remember flag.

    Returns:
        True if the token is valid and not expired, False otherwise.

    Example:
        >>> token = create_session_token()
        >>> verify_session_token(token)
        True
        >>> verify_session_token("invalid_token")
        False
    """
    try:
        # First load without max_age check to read the remember flag
        data = _serializer.loads(token, max_age=None)
        if data.get("admin", False) is not True:
            return False

        # Determine the appropriate max_age based on remember flag
        effective_max_age = (
            REMEMBER_ME_MAX_AGE if data.get("remember", False) else max_age
        )

        # Re-validate with the correct max_age
        data = _serializer.loads(token, max_age=effective_max_age)
        return data.get("admin", False) is True
    except (BadSignature, SignatureExpired):
        return False


def compare_keys(provided_key: str, expected_key: str) -> bool:
    """Compare two API keys in constant time, tolerating any str input.

    secrets.compare_digest raises TypeError when given str arguments that
    contain non-ASCII characters, which turns a bad client key into an
    unhandled 500 instead of a 401. Comparing UTF-8 bytes accepts any
    input while keeping the constant-time guarantee. surrogatepass covers
    lone surrogates, which json.loads can produce from escape sequences
    and which strict UTF-8 encoding rejects.

    Both arguments must be str; None is the caller's responsibility.

    Args:
        provided_key: The key supplied by the client (untrusted).
        expected_key: The configured key to compare against.

    Returns:
        True if the keys match, False otherwise.
    """
    return secrets.compare_digest(
        provided_key.encode("utf-8", "surrogatepass"),
        expected_key.encode("utf-8", "surrogatepass"),
    )


def fingerprint_key(api_key: str) -> str:
    """Return a short, non-reversible fingerprint of an API key for logging.

    Logging a rejected key verbatim leaks the client's secret into the server
    log. A truncated SHA-256 digest lets operators correlate repeated
    rejections of the same key without exposing the key itself. surrogatepass
    matches compare_keys() so any str the auth path accepts can be
    fingerprinted, including lone surrogates from json escape sequences.

    Args:
        api_key: The (untrusted) key to fingerprint. Empty string is allowed.

    Returns:
        The first 8 hex characters of the SHA-256 digest of the UTF-8 bytes.
    """
    digest = hashlib.sha256(api_key.encode("utf-8", "surrogatepass")).hexdigest()
    return digest[:8]


def verify_api_key(api_key: str, server_api_key: str) -> bool:
    """Verify an API key using constant-time comparison.

    This function uses constant-time comparison to prevent timing attacks
    when comparing the provided API key with the server's API key.

    Args:
        api_key: The API key provided by the client.
        server_api_key: The server's configured API key.

    Returns:
        True if the API keys match, False otherwise.

    Example:
        >>> verify_api_key("secret123", "secret123")
        True
        >>> verify_api_key("wrong", "secret123")
        False
    """
    if not api_key or not server_api_key:
        return False
    return compare_keys(api_key, server_api_key)


def verify_any_api_key(api_key: str, main_key: str, sub_keys: list) -> bool:
    """Verify an API key against the main key and all sub keys.

    Uses constant-time comparison for each key to prevent timing attacks.
    Checks the main key first, then iterates through sub keys.

    Args:
        api_key: The API key provided by the client.
        main_key: The server's main API key.
        sub_keys: List of SubKeyEntry objects with .key attribute.

    Returns:
        True if the API key matches any configured key, False otherwise.
    """
    if not api_key:
        return False
    # Check main key
    if main_key and compare_keys(api_key, main_key):
        return True
    # Check sub keys
    for sk in sub_keys:
        if sk.key and compare_keys(api_key, sk.key):
            return True
    return False


def validate_api_key(api_key: str) -> tuple[bool, str]:
    """Validate API key format requirements.

    Rules:
    - Minimum 4 characters
    - No whitespace characters (space, tab, newline, etc.)
    - Printable characters only (no control characters)
    - ASCII characters only

    The ASCII-only rule is not cosmetic: HTTP request headers are decoded as
    latin-1 by the ASGI layer, so a client cannot transmit a non-ASCII key
    intact. A configured key such as "café" therefore starts the server
    fine but can never be matched over the wire, yielding silent 401s on every
    authenticated request. Rejecting it at configuration time surfaces the
    misconfiguration immediately instead.

    Args:
        api_key: The API key string to validate.

    Returns:
        Tuple of (is_valid, error_message). Error message is empty if valid.
    """
    if len(api_key) < 4:
        return False, "API key must be at least 4 characters"
    if any(c.isspace() for c in api_key):
        return False, "API key must not contain whitespace"
    if not api_key.isprintable():
        return False, "API key must contain only printable characters"
    if not api_key.isascii():
        return False, "API key must contain only ASCII characters"
    return True, ""


def verify_session(request: Request) -> bool:
    """Verify if the request has a valid admin session.

    Checks for a valid session cookie in the request.

    Args:
        request: The FastAPI request object.

    Returns:
        True if the session is valid, False otherwise.
    """
    token = request.cookies.get(SESSION_COOKIE_NAME)
    if not token:
        return False
    return verify_session_token(token)


async def require_admin(request: Request) -> bool:
    """FastAPI dependency to require admin authentication.

    This dependency can be used in route definitions to protect
    admin-only endpoints. It checks for a valid session cookie.

    Args:
        request: The FastAPI request object (injected by FastAPI).

    Returns:
        True if authentication is successful.

    Raises:
        HTTPException: 401 Unauthorized if not authenticated.

    Example:
        >>> from fastapi import Depends
        >>> @app.get("/admin/settings")
        ... async def get_settings(is_admin: bool = Depends(require_admin)):
        ...     return {"settings": "..."}
    """
    # Skip admin auth when skip_api_key_verification is enabled
    if _get_global_settings is not None:
        gs = _get_global_settings()
        if gs is not None and gs.auth.skip_api_key_verification:
            return True

    if not verify_session(request):
        # Browser requests (Accept: text/html) get redirected to login page
        accept = request.headers.get("accept", "")
        if "text/html" in accept:
            raise _RedirectToLogin()
        raise HTTPException(
            status_code=401,
            detail="Admin authentication required",
            headers={"WWW-Authenticate": "Cookie"},
        )
    return True


class _RedirectToLogin(Exception):
    """Raised to trigger a redirect to the admin login page."""
    pass


# ===== FILE: omlx/admin/benchmark.py =====

# SPDX-License-Identifier: Apache-2.0
"""Benchmark execution logic for oMLX admin panel.

Provides single-request and continuous-batching benchmarks with
real-time progress reporting via SSE events.
"""

import asyncio
import json
import logging
import os
import re
import time
import uuid
from dataclasses import dataclass, field
from enum import StrEnum
from functools import lru_cache
from pathlib import Path
from typing import Any, Optional

from pydantic import BaseModel, field_validator

from ..utils.proc_memory import get_lifetime_max_phys_footprint
from ..utils.system_sampler import SystemSampler
from .external_api import ExternalAPIClient, ExternalEndpointConfig

try:
    import mlx.core as mx

    HAS_MLX = True
except ImportError:
    HAS_MLX = False

logger = logging.getLogger(__name__)

# Module-level storage for active benchmark runs
_benchmark_runs: dict[str, "BenchmarkRun"] = {}

# Valid prompt lengths for single request tests
VALID_PROMPT_LENGTHS = [1024, 4096, 8192, 16384, 32768, 65536, 131072, 200000]

# Valid batch sizes for continuous batching tests
VALID_BATCH_SIZES = [2, 4, 8]


class BenchmarkContextProfile(StrEnum):
    """Stable identifiers for the bundled throughput-benchmark corpora."""

    CODE_PYTHON = "code_python"
    CODE_MIXED = "code_mixed"
    NOVEL_KO = "novel_ko"
    NOVEL_EN = "novel_en"
    NOVEL_JA = "novel_ja"


@dataclass(frozen=True)
class BenchmarkCorpusSpec:
    """Metadata needed to build local and tokenizer-less prompts."""

    filename: str
    label: str
    chars_per_token: float
    start_marker: str | None = None


BENCHMARK_CONTEXT_PROFILES: dict[BenchmarkContextProfile, BenchmarkCorpusSpec] = {
    BenchmarkContextProfile.CODE_PYTHON: BenchmarkCorpusSpec(
        "code_python.txt", "Code (Python)", 4.0
    ),
    BenchmarkContextProfile.CODE_MIXED: BenchmarkCorpusSpec(
        "code_mixed.txt", "Code (Mixed)", 3.5
    ),
    BenchmarkContextProfile.NOVEL_KO: BenchmarkCorpusSpec(
        "novel_ko.txt", "Novel (Korean)", 1.35
    ),
    BenchmarkContextProfile.NOVEL_EN: BenchmarkCorpusSpec(
        "novel_en.txt", "Novel (English)", 4.0, "Call me Ishmael."
    ),
    BenchmarkContextProfile.NOVEL_JA: BenchmarkCorpusSpec(
        "novel_ja.txt", "Novel (Japanese)", 1.6
    ),
}


class BenchmarkRequest(BaseModel):
    """Request model for starting a benchmark."""

    model_id: str
    prompt_lengths: list[int]
    generation_length: int = 128
    batch_sizes: list[int] = []
    context_profile: BenchmarkContextProfile = BenchmarkContextProfile.CODE_PYTHON
    force_lm_engine: bool = False
    # When set, the benchmark runs against a remote OpenAI-compatible
    # endpoint instead of a local engine and model_id is the remote
    # model name (not validated against the local catalog).
    external: Optional[ExternalEndpointConfig] = None

    @field_validator("prompt_lengths")
    @classmethod
    def validate_prompt_lengths(cls, v: list[int]) -> list[int]:
        if not v:
            raise ValueError("At least one prompt length is required")
        for pl in v:
            if pl not in VALID_PROMPT_LENGTHS:
                raise ValueError(
                    f"Invalid prompt length {pl}. Must be one of {VALID_PROMPT_LENGTHS}"
                )
        return sorted(v)

    @field_validator("batch_sizes")
    @classmethod
    def validate_batch_sizes(cls, v: list[int]) -> list[int]:
        for bs in v:
            if bs not in VALID_BATCH_SIZES:
                raise ValueError(
                    f"Invalid batch size {bs}. Must be one of {VALID_BATCH_SIZES}"
                )
        return sorted(v)


@dataclass
class BenchmarkRun:
    """Tracks the state of a running benchmark.

    SSE delivery model: events are appended to `events` (append-only
    log) under `cond`. Subscribers replay `events` from offset 0 then
    wait on `cond` for new entries. `terminal` is set once the final
    event (`upload_done` / `error`) has been published so subscribers
    know to close their stream rather than wait for a follow-up.
    """

    bench_id: str
    request: BenchmarkRequest
    status: str = "running"  # running, completed, cancelled, error
    events: list[dict] = field(default_factory=list)
    cond: asyncio.Condition = field(default_factory=asyncio.Condition)
    terminal: bool = False
    task: Optional[asyncio.Task] = None
    results: list[dict] = field(default_factory=list)
    error_message: str = ""
    # Acceleration features active when the benchmark started. Results are
    # uploaded either way; the flags ride along so the leaderboard can mark
    # and filter them instead of silently mixing them in.
    experimental_features: list[str] = field(default_factory=list)
    # Same snapshot in the upload payload's shape: [{key, label, detail?}].
    feature_flags: list[dict] = field(default_factory=list)
    # Performance-relevant subset of the model's settings at run start.
    model_settings_snapshot: Optional[dict] = None
    # Host telemetry sampler, running for the duration of the tests.
    sampler: Optional[Any] = None
    # Lifetime footprint high-water mark before the tests began, so the run's
    # own peak can be told apart from a larger one set earlier in the process.
    lifetime_footprint_at_start: int = 0
    # Mirror of the upload SSE events so REST consumers (e.g. native Swift
    # app polling /results) can render leaderboard status without opening
    # the stream. Phases: "idle" → "uploading" → "done" | "skipped". The
    # browser HTML still consumes the SSE stream directly; this is purely
    # additive state that lives alongside it.
    upload_state: dict = field(
        default_factory=lambda: {
            "phase": "idle",
            "results": [],  # per-context-length: {context_length, id?, url?, duplicate?, error?}
            "total": 0,
            "success_count": 0,
            "failed_count": 0,
            "owner_hash": None,  # display hash, populated on upload_done
            "skipped_reason": None,  # only "external_endpoint" reaches this now
            # Always empty. Kept because BenchDTO.swift declares it non-optional,
            # so dropping the key would fail decoding on every app build that has
            # not been updated — which turns into a per-second error loop while the
            # results poller runs.
            "skipped_features": [],
            "feature_flags": [],  # [{key, label, detail?}]
        }
    )


# Event types that close the SSE stream for a bench run. `done` is NOT
# terminal — it marks "tests finished, upload starting"; the real end of
# stream is `upload_done` (or `error`). `upload_skipped` is the external
# endpoint's last event: without it here, subscribers to an external run would
# wait for an `upload_done` that never comes.
_BENCH_TERMINAL_TYPES = frozenset({"upload_done", "upload_skipped", "error"})


@dataclass(frozen=True)
class _FeatureFlagSpec:
    """One acceleration toggle, in both the legacy and upload projections."""

    attr: str
    legacy: str
    key: str
    label: str
    detail_attr: Optional[str] = None
    detail_key_fmt: Optional[str] = None
    detail_label_fmt: Optional[str] = None


# `mtp_enabled` is surfaced as "Lightning MTP" everywhere in the UI, so the
# upload key follows the user-facing name rather than the settings field.
_FEATURE_FLAG_SPECS = (
    _FeatureFlagSpec("dflash_enabled", "dflash", "dflash", "DFlash"),
    _FeatureFlagSpec(
        "specprefill_enabled", "specprefill", "specprefill", "SpecPrefill"
    ),
    _FeatureFlagSpec(
        "turboquant_kv_enabled",
        "turboquant",
        "turboquant_kv",
        "TurboQuant KV",
        detail_attr="turboquant_kv_bits",
        detail_key_fmt="_{}bit",
        detail_label_fmt=" {}-bit",
    ),
    _FeatureFlagSpec("mtp_enabled", "mtp", "lightning_mtp", "Lightning MTP"),
    _FeatureFlagSpec("vlm_mtp_enabled", "vlm_mtp", "vlm_mtp", "VLM MTP"),
)


def _sample_window(run: "BenchmarkRun", window_start: float) -> Optional[dict]:
    """Aggregate host telemetry for the interval a single test occupied."""
    if run.sampler is None:
        return None
    try:
        return run.sampler.window(window_start, time.monotonic())
    except Exception as e:  # noqa: BLE001
        logger.debug(f"Benchmark: system metrics unavailable: {e}")
        return None


def _detect_experimental_features(model_settings: Any) -> list[str]:
    """Return benchmark-skewing model features enabled in settings."""
    return [
        spec.legacy
        for spec in _FEATURE_FLAG_SPECS
        if getattr(model_settings, spec.attr, False)
    ]


def _format_bits(value: Any) -> Optional[str]:
    """Render a bit-width for display, dropping a trailing .0 (4.0 -> "4")."""
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return f"{number:g}"


def _derive_feature_flags(model_settings: Any) -> list[dict]:
    """Build the upload projection of the active acceleration features.

    Objects rather than bare keys: the app and omlx.ai ship independently, so
    carrying the display label means a newly added feature renders correctly on
    the site from day one instead of showing a raw snake_case key until the
    next site deploy. Only active features are included — the site derives
    "this run was accelerated" from the list being non-empty.
    """
    flags: list[dict] = []
    for spec in _FEATURE_FLAG_SPECS:
        if not getattr(model_settings, spec.attr, False):
            continue
        key, label = spec.key, spec.label
        if spec.detail_attr:
            bits = _format_bits(getattr(model_settings, spec.detail_attr, None))
            if bits:
                # Keys must stay [a-z0-9_], so 2.5 becomes 2_5.
                key += spec.detail_key_fmt.format(bits.replace(".", "_"))
                label += spec.detail_label_fmt.format(bits)
        flags.append({"key": key, "label": label})
    return flags


# Performance-relevant settings only, as an allowlist rather than a denylist:
# ModelSettings gains fields regularly, and a denylist would ship every future
# addition to a public endpoint by default.
#
# Excluded on purpose: display_name / description / model_alias (user-authored
# free text), is_pinned / is_default / is_hidden / is_favorite /
# active_profile_name / ttl_seconds (local organization), the guided_grammar
# body (unbounded; the boolean is kept), chat_template_kwargs and
# forced_ct_kwargs (arbitrary user dicts), and trust_remote_code (security
# posture, not performance). The *_draft_model fields are included but reduced
# to a basename — the drafter's identity explains an MTP/DFlash result, while
# the full path would leak the local filesystem layout and the OS username.
_UPLOADED_SETTING_FIELDS = (
    "max_context_window",
    "max_tokens",
    "temperature",
    "top_p",
    "top_k",
    "min_p",
    "repetition_penalty",
    "presence_penalty",
    "force_sampling",
    "enable_thinking",
    "thinking_budget_enabled",
    "thinking_budget_tokens",
    "reasoning_parser",
    "guided_grammar_enabled",
    "model_type_override",
    "index_cache_freq",
    "turboquant_kv_enabled",
    "turboquant_kv_bits",
    "turboquant_skip_last",
    "specprefill_enabled",
    "specprefill_draft_model",
    "specprefill_keep_pct",
    "specprefill_threshold",
    "dflash_enabled",
    "dflash_draft_model",
    "dflash_draft_quant_enabled",
    "dflash_draft_quant_weight_bits",
    "dflash_draft_quant_activation_bits",
    "dflash_draft_quant_group_size",
    "dflash_max_ctx",
    "dflash_in_memory_cache",
    "dflash_in_memory_cache_max_entries",
    "dflash_ssd_cache",
    "dflash_draft_window_size",
    "dflash_draft_sink_size",
    "dflash_verify_mode",
    "mtp_enabled",
    "mtp_num_draft_tokens",
    "vlm_mtp_enabled",
    "vlm_mtp_draft_model",
    "vlm_mtp_draft_block_size",
)

_PATH_VALUED_SETTING_FIELDS = frozenset(
    {
        "specprefill_draft_model",
        "dflash_draft_model",
        "vlm_mtp_draft_model",
    }
)

_MAX_UPLOADED_SETTINGS_BYTES = 4096


def _filter_uploaded_settings(model_settings: Any) -> Optional[dict]:
    """Project model settings onto the uploadable allowlist."""
    to_dict = getattr(model_settings, "to_dict", None)
    if not callable(to_dict):
        return None
    try:
        raw = to_dict()
    except Exception as e:  # noqa: BLE001
        logger.warning(f"Benchmark: failed to serialize model settings: {e}")
        return None

    filtered: dict = {}
    for key in _UPLOADED_SETTING_FIELDS:
        if key not in raw:
            continue
        value = raw[key]
        if key in _PATH_VALUED_SETTING_FIELDS and isinstance(value, str):
            value = os.path.basename(value.rstrip("/")) or value
        filtered[key] = value

    if len(json.dumps(filtered, separators=(",", ":"))) > _MAX_UPLOADED_SETTINGS_BYTES:
        logger.warning(
            "Benchmark: model settings snapshot exceeded "
            f"{_MAX_UPLOADED_SETTINGS_BYTES} bytes, uploading accelerator flags only"
        )
        filtered = {
            spec.attr: filtered[spec.attr]
            for spec in _FEATURE_FLAG_SPECS
            if spec.attr in filtered
        }
    return filtered


def _with_benchmark_context(
    context_profile: BenchmarkContextProfile | str,
    model_settings: dict | None,
) -> dict:
    """Prepend the benchmark context to the uploaded settings snapshot."""
    settings = dict(model_settings or {})
    settings.pop("benchmark_context", None)
    return {
        "benchmark_context": benchmark_context_label(context_profile),
        **settings,
    }


def get_run(bench_id: str) -> Optional[BenchmarkRun]:
    """Get a benchmark run by ID."""
    return _benchmark_runs.get(bench_id)


def get_active_run() -> Optional[BenchmarkRun]:
    """Return the currently-running throughput benchmark, if any.

    Discovery surface for clients that need to attach to an in-progress
    run without knowing the bench_id upfront (page refresh, second tab).
    Returns the first run with status == "running"; throughput benches
    are 1-at-a-time so there's never more than one.
    """
    for run in _benchmark_runs.values():
        if run.status == "running":
            return run
    return None


def create_run(request: BenchmarkRequest) -> BenchmarkRun:
    """Create and register a new benchmark run."""
    bench_id = f"bench-{uuid.uuid4().hex[:12]}"
    run = BenchmarkRun(bench_id=bench_id, request=request)
    _benchmark_runs[bench_id] = run
    return run


def cleanup_old_runs(max_runs: int = 10) -> None:
    """Remove old completed runs to prevent memory leaks."""
    completed = [
        (bid, r)
        for bid, r in _benchmark_runs.items()
        if r.status in ("completed", "cancelled", "error")
    ]
    if len(completed) > max_runs:
        for bid, _ in completed[:-max_runs]:
            del _benchmark_runs[bid]


# Bundled corpora for benchmark prompts. They contain long-form code or prose,
# never a short filler sentence. A whole corpus may repeat for a tokenizer that
# compresses it unusually well, but the repeated unit is hundreds of thousands
# of natural tokens rather than a predictable one-line loop.
_BENCH_CORPUS_DIR = Path(__file__).parent / "bench_corpora"
_PROMPT_BUILD_MAX_ATTEMPTS = 16


def benchmark_context_label(profile: BenchmarkContextProfile | str) -> str:
    """Return the user-facing label for a benchmark context profile."""
    normalized = BenchmarkContextProfile(profile)
    return BENCHMARK_CONTEXT_PROFILES[normalized].label


@lru_cache(maxsize=len(BENCHMARK_CONTEXT_PROFILES))
def _load_bench_corpus(
    context_profile: (
        BenchmarkContextProfile | str
    ) = BenchmarkContextProfile.CODE_PYTHON,
) -> str:
    profile = BenchmarkContextProfile(context_profile)
    spec = BENCHMARK_CONTEXT_PROFILES[profile]
    path = _BENCH_CORPUS_DIR / spec.filename
    corpus = path.read_text(encoding="utf-8")
    if spec.start_marker:
        start = corpus.find(spec.start_marker)
        if start < 0:
            raise RuntimeError(
                f"Benchmark corpus at {path} is missing the content start marker"
            )
        corpus = corpus[start:]
    if not corpus:
        raise RuntimeError(f"Benchmark corpus at {path} is empty")
    return corpus


def _generate_prompt(
    tokenizer: Any,
    target_tokens: int,
    context_profile: (
        BenchmarkContextProfile | str
    ) = BenchmarkContextProfile.CODE_PYTHON,
) -> list[int]:
    """Generate exactly ``target_tokens`` benchmark-corpus token IDs.

    Uses a unique UUID prefix to prevent SSD cache hits from previous sessions.
    The prefix and corpus are encoded together so tokenizer boundary merges and
    special-token insertion happen exactly once. The token IDs are passed to the
    engine directly; decoding them back to text would make exact length depend
    on tokenizer round-trip behavior.
    """
    if target_tokens <= 0:
        raise ValueError("target_tokens must be positive")

    unique_prefix = f"BENCH-{uuid.uuid4().hex} "
    profile = BenchmarkContextProfile(context_profile)
    spec = BENCHMARK_CONTEXT_PROFILES[profile]
    corpus = _load_bench_corpus(profile)

    target_chars = max(round(target_tokens * spec.chars_per_token), 1)
    for _ in range(_PROMPT_BUILD_MAX_ATTEMPTS):
        repeats = (target_chars + len(corpus) - 1) // len(corpus)
        body = (corpus * repeats)[:target_chars]
        tokens = [int(token) for token in tokenizer.encode(unique_prefix + body)]
        if len(tokens) >= target_tokens:
            return tokens[:target_tokens]
        if not tokens:
            raise RuntimeError(
                f"Benchmark corpus {profile.value} tokenized to 0 tokens"
            )

        # Scale by the observed tokenizer ratio, rounding up. The +1 guarantees
        # progress even for tokenizers with coarse or unusual segmentation.
        target_chars = max(
            target_chars + 1,
            (target_chars * target_tokens + len(tokens) - 1) // len(tokens) + 1,
        )

    raise RuntimeError(
        f"Could not build an exact {target_tokens}-token benchmark prompt "
        f"after {_PROMPT_BUILD_MAX_ATTEMPTS} attempts"
    )


def _generate_external_prompt(
    target_tokens: int,
    context_profile: (
        BenchmarkContextProfile | str
    ) = BenchmarkContextProfile.CODE_PYTHON,
) -> str:
    """Generate an approximately target_tokens-long prompt without a tokenizer.

    Uses a unique UUID prefix so remote prefix caches cannot skew results.
    """
    unique_prefix = f"BENCH-{uuid.uuid4().hex} "
    profile = BenchmarkContextProfile(context_profile)
    spec = BENCHMARK_CONTEXT_PROFILES[profile]
    corpus = _load_bench_corpus(profile)
    target_chars = max(
        0,
        round(target_tokens * spec.chars_per_token) - len(unique_prefix),
    )
    repeats = (target_chars + len(corpus) - 1) // len(corpus)
    return unique_prefix + (corpus * repeats)[:target_chars]


def _compute_single_metrics(
    prompt_tokens: int,
    completion_tokens: int,
    start_time: float,
    first_token_time: float,
    end_time: float,
    peak_memory: int,
    cached_tokens: int,
    prefill_duration_s: float | None = None,
    generation_duration_s: float | None = None,
    generation_measured: bool = True,
    timing_observed: bool = True,
) -> dict:
    """Compute all metrics for a single request benchmark."""
    ttft_s = first_token_time - start_time
    prefill_duration = prefill_duration_s if prefill_duration_s is not None else ttft_s
    gen_duration = (
        generation_duration_s
        if generation_duration_s is not None
        else end_time - first_token_time
    )
    e2e_duration = end_time - start_time

    ttft_ms: float | None = ttft_s * 1000
    if generation_measured and completion_tokens > 1 and gen_duration > 0:
        tpot_ms: float | None = (gen_duration / (completion_tokens - 1)) * 1000
        gen_tps: float | None = completion_tokens / gen_duration
    else:
        # Generation timing could not be measured (e.g. all content arrived
        # in a single burst with no measurable inter-token span) — report
        # unmeasured rather than a misleading 0.0.
        tpot_ms = None
        gen_tps = None
    processing_tps: float | None = prompt_tokens / max(prefill_duration, 1e-9)
    total_throughput = (prompt_tokens + completion_tokens) / max(e2e_duration, 1e-9)

    if not timing_observed:
        # The first-token timestamp was never observed and fell back to the
        # end of the response, so TTFT covers the whole response and the
        # prefill rate derived from it is not a prefill rate at all. Only
        # e2e latency and total throughput survive.
        ttft_ms = None
        processing_tps = None

    return {
        "ttft_ms": round(ttft_ms, 1) if ttft_ms is not None else None,
        "tpot_ms": round(tpot_ms, 2) if tpot_ms is not None else None,
        "gen_tps": round(gen_tps, 1) if gen_tps is not None else None,
        "processing_tps": (
            round(processing_tps, 1) if processing_tps is not None else None
        ),
        "e2e_latency_s": round(e2e_duration, 3),
        "total_throughput": round(total_throughput, 1),
        "peak_memory_bytes": peak_memory,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "cached_tokens": cached_tokens,
    }


def _pin_speed_priority(engine_pool: Any) -> bool | None:
    """Force prefill speed priority for the benchmark's duration.

    Throughput numbers measured while the memory-guard throttle shrinks
    prefill chunks are not comparable, so the bench always runs in speed
    mode. The pin lands on the pool's stored scheduler config — the bench
    model is loaded fresh after unload-all, so that is where its Scheduler
    reads the flag from. Returns the previous value for restoration, or
    None when the pool exposes no config (nothing to restore).
    """
    config = getattr(engine_pool, "_scheduler_config", None)
    if config is None:
        return None
    previous = bool(getattr(config, "prefill_speed_priority", False))
    config.prefill_speed_priority = True
    return previous


def _restore_speed_priority(engine_pool: Any, previous: bool | None) -> None:
    """Undo _pin_speed_priority (no-op when the pin never landed)."""
    if previous is None:
        return
    config = getattr(engine_pool, "_scheduler_config", None)
    if config is not None:
        config.prefill_speed_priority = previous


def _get_batch_benchmark_core(engine: Any) -> Any | None:
    """Return the scheduler core when this engine supports batch benchmarks."""
    engine_core = getattr(engine, "_engine", None)
    if engine_core is None:
        return None
    if not callable(getattr(engine_core, "add_request", None)):
        return None
    if not callable(getattr(engine_core, "stream_outputs", None)):
        return None
    return engine_core


async def _send_event(run: BenchmarkRun, event: dict) -> None:
    """Append an event to the run's log and wake any subscribers.

    Sets `run.terminal` when the event ends the stream so subscribers
    can return rather than wait for an event that will never come.
    """
    async with run.cond:
        run.events.append(event)
        if event.get("type") in _BENCH_TERMINAL_TYPES:
            run.terminal = True
        run.cond.notify_all()


async def _run_single_test(
    engine: Any,
    prompt: list[int],
    max_tokens: int,
    pp_len: int,
) -> dict:
    """Run a single request benchmark test and return metrics."""
    if len(prompt) != pp_len:
        raise RuntimeError(
            f"Benchmark prompt length mismatch before pp{pp_len}: "
            f"built {len(prompt)} tokens"
        )

    # Reset peak memory tracking
    try:
        mx.reset_peak_memory()
    except Exception:
        pass

    start_time = time.perf_counter()
    first_token_time = None
    last_generated_token_time = None
    last_output = None
    prev_completion_tokens = 0

    async for output in engine.stream_generate(
        prompt=prompt,
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=1.0,
    ):
        # Detect first generated token via completion_tokens count,
        # not new_text. Some models (e.g. Harmony/gpt-oss) produce
        # protocol tokens that don't yield visible new_text.
        completion_delta = output.completion_tokens - prev_completion_tokens
        if completion_delta > 0:
            generated_at = getattr(output, "generated_at", None)
            generated_until = getattr(output, "generated_until", None)
            output_first_token_time = (
                float(generated_at) if generated_at is not None else time.perf_counter()
            )
            if first_token_time is None:
                first_token_time = output_first_token_time
            if generated_until is not None:
                last_generated_token_time = float(generated_until)
            elif completion_delta == 1:
                last_generated_token_time = output_first_token_time
        prev_completion_tokens = output.completion_tokens
        last_output = output

    end_time = time.perf_counter()

    if first_token_time is None:
        first_token_time = end_time

    # Get peak memory
    try:
        peak_memory = mx.get_peak_memory()
    except Exception:
        peak_memory = 0

    if last_output is None:
        raise RuntimeError(f"Benchmark pp{pp_len} produced no engine output")

    prompt_tokens = last_output.prompt_tokens
    if prompt_tokens != pp_len:
        raise RuntimeError(
            f"Benchmark prompt length mismatch after pp{pp_len}: "
            f"engine reported {prompt_tokens} tokens"
        )

    completion_tokens = last_output.completion_tokens
    cached_tokens = last_output.cached_tokens

    if cached_tokens > 0:
        logger.warning(
            f"Benchmark test pp{pp_len} had {cached_tokens} cached tokens "
            f"(expected 0). Results may not reflect true prefill performance."
        )

    prefill_duration_s = None
    generation_duration_s = None
    producer_generation_duration_s = None
    metric_completion_tokens = completion_tokens
    if first_token_time is not None and last_generated_token_time is not None:
        measured_duration = last_generated_token_time - first_token_time
        if measured_duration > 0:
            producer_generation_duration_s = measured_duration
    if last_output is not None:
        prompt_tps = float(getattr(last_output, "prompt_tps", 0.0) or 0.0)
        if prompt_tps > 0 and prompt_tokens > 0:
            prefill_duration_s = prompt_tokens / prompt_tps

        canvas_tps = float(getattr(last_output, "diffusion_canvas_tps", 0.0) or 0.0)
        canvas_tokens = int(getattr(last_output, "diffusion_canvas_tokens", 0) or 0)
        if canvas_tps > 0 and canvas_tokens > 0:
            metric_completion_tokens = canvas_tokens
            generation_duration_s = canvas_tokens / canvas_tps
        else:
            generation_tps = float(getattr(last_output, "generation_tps", 0.0) or 0.0)
            if generation_tps > 0 and completion_tokens > 0:
                generation_duration_s = completion_tokens / generation_tps

    if generation_duration_s is None:
        generation_duration_s = producer_generation_duration_s

    generation_measured = generation_duration_s is not None

    return _compute_single_metrics(
        prompt_tokens=prompt_tokens,
        completion_tokens=metric_completion_tokens,
        start_time=start_time,
        first_token_time=first_token_time,
        end_time=end_time,
        peak_memory=peak_memory,
        cached_tokens=cached_tokens,
        prefill_duration_s=prefill_duration_s,
        generation_duration_s=generation_duration_s,
        generation_measured=generation_measured,
    )


async def _run_batch_test(
    engine: Any,
    prompts: list[list[int]],
    prompt_tokens: int,
    max_tokens: int,
    batch_size: int,
) -> dict:
    """Run a continuous batching benchmark test.

    Submits batch_size concurrent requests via the engine core and measures
    aggregate throughput including pp TPS and tg TPS.

    Args:
        prompts: List of prompts (one per request). For same-prompt tests,
                 all entries are identical. For different-prompt tests, each
                 has a unique UUID prefix.
        prompt_tokens: Number of prompt tokens per request (for pp TPS calc).
    """
    from ..request import SamplingParams

    engine_core = _get_batch_benchmark_core(engine)
    if engine_core is None:
        raise ValueError("Engine does not support batch benchmarks")
    if len(prompts) < batch_size:
        raise RuntimeError(
            f"Benchmark batch requires {batch_size} prompts, got {len(prompts)}"
        )
    invalid_lengths = [
        len(prompt) for prompt in prompts[:batch_size] if len(prompt) != prompt_tokens
    ]
    if invalid_lengths:
        raise RuntimeError(
            f"Benchmark batch prompt length mismatch: expected {prompt_tokens}, "
            f"got {invalid_lengths}"
        )

    sampling_params = SamplingParams(
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=1.0,
    )

    async def _single_request(prompt: list[int]) -> dict:
        """Run a single request within the batch."""
        start = time.perf_counter()
        first_token = None
        tokens = 0
        prev_tokens = 0
        reported_prompt_tokens = 0

        request_id = await engine_core.add_request(
            prompt=prompt,
            sampling_params=sampling_params,
        )

        async for output in engine_core.stream_outputs(request_id):
            if first_token is None and output.completion_tokens > prev_tokens:
                first_token = time.perf_counter()
            prev_tokens = output.completion_tokens
            if output.finished:
                tokens = output.completion_tokens
                reported_prompt_tokens = output.prompt_tokens

        end = time.perf_counter()
        if first_token is None:
            first_token = end
        if reported_prompt_tokens != prompt_tokens:
            raise RuntimeError(
                f"Benchmark batch prompt length mismatch after submission: "
                f"expected {prompt_tokens}, engine reported "
                f"{reported_prompt_tokens}"
            )

        return {
            "ttft_s": first_token - start,
            "first_token_abs": first_token,
            "end_abs": end,
            "completion_tokens": tokens,
        }

    # Submit all requests concurrently. Nothing else generates during the
    # gather, so the process-global peak counter belongs to this batch.
    if HAS_MLX:
        try:
            mx.reset_peak_memory()
        except Exception:
            pass
    wall_start = time.perf_counter()
    results = await asyncio.gather(
        *[_single_request(prompts[i]) for i in range(batch_size)]
    )
    wall_end = time.perf_counter()
    peak_memory = 0
    if HAS_MLX:
        try:
            peak_memory = mx.get_peak_memory()
        except Exception:
            peak_memory = 0

    # Aggregate metrics
    total_gen_tokens = sum(r["completion_tokens"] for r in results)
    total_prompt_tokens = prompt_tokens * batch_size
    wall_time = wall_end - wall_start
    avg_ttft_ms = (sum(r["ttft_s"] for r in results) / batch_size) * 1000

    # pp TPS: total prompt tokens / time until ALL requests finish prefill
    max_first_token = max(r["first_token_abs"] for r in results)
    prefill_wall_time = max_first_token - wall_start
    pp_tps = total_prompt_tokens / max(prefill_wall_time, 1e-9)

    # tg TPS: total generated tokens / generation wall time
    # Generation starts when the last request finishes prefill
    gen_wall_time = wall_end - max_first_token
    tg_tps = total_gen_tokens / max(gen_wall_time, 1e-9)

    return {
        "pp_tps": round(pp_tps, 1),
        "tg_tps": round(tg_tps, 1),
        "avg_ttft_ms": round(avg_ttft_ms, 1),
        "e2e_latency_s": round(wall_time, 3),
        "peak_memory_bytes": peak_memory,
        "total_gen_tokens": total_gen_tokens,
        "batch_size": batch_size,
    }


async def _run_external_single_test(
    client: ExternalAPIClient,
    prompt: str,
    max_tokens: int,
) -> dict:
    """Run a single-request benchmark against an external endpoint.

    Token counts come from the endpoint's streamed usage payload, never
    from counting SSE chunks (providers batch multiple tokens per chunk).
    Prefill duration is not observable remotely, so pp TPS falls back to
    prompt_tokens / TTFT (network latency included). Peak memory is not
    measurable for a remote host.
    """
    stats = await client.stream_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=0.0,
    )
    gen_duration = stats.last_content_time - stats.first_content_time
    metrics = _compute_single_metrics(
        prompt_tokens=stats.prompt_tokens,
        completion_tokens=stats.completion_tokens,
        start_time=stats.start_time,
        first_token_time=stats.first_content_time,
        end_time=stats.end_time,
        peak_memory=0,
        cached_tokens=stats.cached_tokens,
        prefill_duration_s=None,
        generation_duration_s=gen_duration if gen_duration > 0 else None,
        generation_measured=gen_duration > 0,
        timing_observed=stats.content_observed,
    )
    metrics["peak_memory_bytes"] = None
    return metrics


async def _run_external_batch_test(
    client: ExternalAPIClient,
    prompts: list[str],
    max_tokens: int,
    batch_size: int,
) -> dict:
    """Run a concurrent-requests benchmark against an external endpoint.

    Mirrors _run_batch_test aggregation, with actual per-request token
    counts taken from each stream's usage payload.
    """
    wall_start = time.perf_counter()
    stats_list = await asyncio.gather(
        *[
            client.stream_chat_completion(
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                temperature=0.0,
            )
            for prompt in prompts
        ]
    )
    wall_end = time.perf_counter()

    total_gen_tokens = sum(s.completion_tokens for s in stats_list)
    prompt_tokens_per_request = [s.prompt_tokens for s in stats_list]
    total_prompt_tokens = sum(prompt_tokens_per_request)
    wall_time = wall_end - wall_start

    # Every aggregate below is derived from the per-request content
    # timestamps, so a single stream that never reported content poisons all
    # of them: its fallback timestamp sits at the end of the response and
    # drags max_first_token along with it.
    timing_observed = all(s.content_observed for s in stats_list)
    decode_observed = all(
        s.last_content_time > s.first_content_time for s in stats_list
    )

    max_first_token = max(s.first_content_time for s in stats_list)
    gen_window = wall_end - max_first_token

    avg_ttft_ms: float | None = None
    pp_tps: float | None = None
    tg_tps: float | None = None
    if timing_observed:
        total_ttft_s = sum(s.first_content_time - s.start_time for s in stats_list)
        avg_ttft_ms = round((total_ttft_s / batch_size) * 1000, 1)
        # pp TPS: total prompt tokens / time until ALL requests emit content
        prefill_window = max(max_first_token - wall_start, 1e-9)
        pp_tps = round(total_prompt_tokens / prefill_window, 1)
        # tg TPS needs a real decode span. wall_end is sampled after
        # asyncio.gather returns, so gen_window stays positive even when
        # every per-request timestamp collapsed onto the end of the
        # response, and the window alone cannot tell a genuine decode phase
        # from a single-chunk dump.
        if decode_observed and gen_window > 0:
            tg_tps = round(total_gen_tokens / gen_window, 1)

    return {
        "pp_tps": pp_tps,
        "tg_tps": tg_tps,
        "avg_ttft_ms": avg_ttft_ms,
        "e2e_latency_s": round(wall_time, 3),
        "total_gen_tokens": total_gen_tokens,
        "total_prompt_tokens": total_prompt_tokens,
        "prompt_tokens": round(total_prompt_tokens / batch_size),
        "prompt_tokens_min": min(prompt_tokens_per_request),
        "prompt_tokens_max": max(prompt_tokens_per_request),
        "batch_size": batch_size,
    }


OMLX_AI_API_URL = "https://omlx.ai/api/benchmarks"

# The leaderboard accepts model_name up to 150 characters.
_MAX_MODEL_NAME_LEN = 150


def _detect_quantization(model_path: str) -> str:
    """Detect model quantization from config.json or directory name.

    Fallback chain: config.json → directory name → "unknown"
    """
    config_path = Path(model_path) / "config.json"
    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
            qconfig = config.get("quantization_config", {})
            bits = qconfig.get("bits")
            if bits is not None:
                return f"{bits}bit"
        except Exception:
            pass

    # Fallback: extract from directory name
    dirname = Path(model_path).name
    match = re.search(
        r"(2bit|3bit|4bit|6bit|8bit|fp16|bf16|MXFP4|NVFP4)", dirname, re.IGNORECASE
    )
    if match:
        return match.group(1).lower()

    return "unknown"


def _upload_model_name(model_id: str) -> str:
    """Model name to publish: exactly what oMLX shows and its copy button copies.

    Quantization and MLX suffixes used to be stripped here, which lost the one
    detail that distinguishes two builds of the same model on the leaderboard.
    The trailing path component is taken defensively — discovery registers ids
    as a single path component today, so this is a no-op for local runs.
    """
    name = model_id.rstrip("/").split("/")[-1]
    return name[:_MAX_MODEL_NAME_LEN]


def _sanitize_upload_error(resp: Any) -> str:
    """Extract a user-presentable error string from a failed upload response.

    Avoids dumping raw HTML bodies (e.g. Cloudflare's "Just a moment..."
    challenge interstitial) into the dashboard's red-x error column.
    Detects CF mitigation specifically so users get actionable context
    instead of a 5KB markup blob.

    Resolution order:
    1. Cloudflare challenge — header ``cf-mitigated: challenge`` is
       authoritative; a body sniff for "just a moment" / "cf-chl" covers
       edge transports that strip the header.
    2. JSON envelope — the omlx.ai API's normal error shape; extract
       ``error`` / ``detail`` / ``message`` if present, truncated.
    3. Plain-text body — short responses only; HTML-looking bodies are
       collapsed to a one-line "non-JSON response (N bytes)" hint.
    4. Fallback to the bare HTTP status code.
    """
    headers = getattr(resp, "headers", {}) or {}
    cf_mitigated = str(headers.get("cf-mitigated", "")).lower()
    body = getattr(resp, "text", "") or ""
    status = getattr(resp, "status_code", "?")

    body_head = body[:512].lower()
    if (
        cf_mitigated == "challenge"
        or "just a moment" in body_head
        or "cf-chl" in body_head
    ):
        return (
            f"Upload blocked by Cloudflare (HTTP {status}). "
            f"This is a server-side issue with omlx.ai — retry later or "
            f"report it to the maintainer."
        )

    try:
        data = resp.json()
        msg = data.get("error") or data.get("detail") or data.get("message")
        if msg:
            return str(msg)[:300]
    except Exception:
        pass

    text = body.strip()
    if "<" in text and ">" in text:
        return f"HTTP {status} — unexpected non-JSON response ({len(body)} bytes)"
    return text[:300] or f"HTTP {status}"


async def _upload_to_omlx_ai(run: BenchmarkRun, engine_pool: Any) -> None:
    """Upload benchmark results to omlx.ai community benchmarks.

    Sends each single-request result as a separate submission,
    grouped by submission_group. Upload failures don't affect
    the benchmark run status.
    """
    import requests

    from .._version import __version__
    from ..utils.hardware import (
        compute_owner_hash,
        get_chip_name,
        get_gpu_core_count,
        get_io_platform_uuid,
        get_os_version,
        get_total_memory_gb,
        parse_chip_info,
    )

    # Accelerated runs upload too. They carry their flags so the leaderboard
    # can mark and filter them, which is more useful than withholding the one
    # set of numbers people most want to see.
    run.upload_state["feature_flags"] = list(run.feature_flags)
    if run.feature_flags:
        logger.info(
            "Benchmark upload tagged with acceleration flags: "
            f"{[f['key'] for f in run.feature_flags]}"
        )

    run.upload_state["phase"] = "uploading"
    await _send_event(
        run,
        {
            "type": "progress",
            "phase": "upload",
            "message": "Uploading to community benchmarks...",
            "current": 0,
            "total": 0,
        },
    )

    # Collect hardware info
    chip_string = get_chip_name()
    chip_name, chip_variant = parse_chip_info(chip_string)
    memory_gb = round(get_total_memory_gb())
    gpu_cores = get_gpu_core_count()
    os_version = get_os_version()
    omlx_version = __version__

    # Compute owner_hash
    owner_hash_full = None
    owner_hash_display = None
    io_uuid = get_io_platform_uuid()
    if io_uuid:
        owner_hash_full = compute_owner_hash(io_uuid, chip_name, gpu_cores, memory_gb)
        # Display hash is without the verify character
        owner_hash_display = owner_hash_full[:-1]

    # Get model info
    entry = engine_pool.get_entry(run.request.model_id)
    model_path = entry.model_path if entry else ""
    quantization = _detect_quantization(model_path)
    model_name = _upload_model_name(run.request.model_id)

    # Generate submission group
    submission_group = str(uuid.uuid4())

    # Peak process memory for the run. ri_lifetime_max_phys_footprint is a
    # high-water mark since process start, so it only describes this benchmark
    # when the benchmark actually set a new maximum — a server that previously
    # held a larger model would otherwise report that older peak. Fall back to
    # the sampler's own maximum, which is scoped to the run.
    peak_footprint_gb = None
    lifetime_end = get_lifetime_max_phys_footprint()
    peak_bytes = 0
    if lifetime_end and lifetime_end > run.lifetime_footprint_at_start:
        peak_bytes = lifetime_end
    elif run.sampler is not None:
        peak_bytes = run.sampler.run_peak_footprint()
    if peak_bytes > 0:
        peak_footprint_gb = round(peak_bytes / (1024**3), 2)

    # Collect single results and batch results
    single_results = [r for r in run.results if r.get("test_type") == "single"]
    uploadable_single_results = [
        r for r in single_results if float(r.get("gen_tps", 0.0) or 0.0) > 0.0
    ]
    skipped_count = len(single_results) - len(uploadable_single_results)
    batch_results = [r for r in run.results if r.get("test_type") == "batch"]

    # Build batching_results from batch data
    batching_results = []
    pp1024_single = next((r for r in single_results if r.get("pp") == 1024), None)
    if (
        pp1024_single
        and float(pp1024_single.get("gen_tps", 0.0) or 0.0) > 0.0
        and batch_results
    ):
        baseline_tps = pp1024_single["gen_tps"]
        batching_results.append(
            {
                "batch_size": 1,
                "tg_tps": baseline_tps,
                "speedup": 1.0,
            }
        )
        for br in batch_results:
            speedup = round(br["tg_tps"] / baseline_tps, 2) if baseline_tps > 0 else 1.0
            batching_results.append(
                {
                    "batch_size": br["batch_size"],
                    "tg_tps": br["tg_tps"],
                    "speedup": speedup,
                }
            )

    success_count = 0
    failed_count = 0

    if skipped_count:
        logger.info(
            f"Benchmark upload skipped {skipped_count} result(s) without "
            f"measurable generation throughput"
        )

    for result in uploadable_single_results:
        context_length = result["pp"]
        peak_mem_gb = None
        if result.get("peak_memory_bytes") and result["peak_memory_bytes"] > 0:
            peak_mem_gb = round(result["peak_memory_bytes"] / (1024**3), 2)

        payload = {
            "chip_name": chip_name,
            "chip_variant": chip_variant,
            "memory_gb": memory_gb,
            "gpu_cores": gpu_cores,
            "omlx_version": omlx_version,
            "os_version": os_version,
            "model_name": model_name,
            "quantization": quantization,
            "context_length": context_length,
            "context_profile": run.request.context_profile.value,
            "pp_tps": result["processing_tps"],
            "tg_tps": result["gen_tps"],
            "ttft_ms": result.get("ttft_ms"),
            "peak_memory_gb": peak_mem_gb,
            "submission_group": submission_group,
            "peak_footprint_gb": peak_footprint_gb,
            "feature_flags": run.feature_flags,
            "model_settings": _with_benchmark_context(
                run.request.context_profile,
                run.model_settings_snapshot,
            ),
            # Per-row: each context length has its own load window. Stays None
            # when sampling was unavailable, so the site does not average
            # fabricated zeros in as measurements.
            "system_metrics": result.get("system_metrics"),
        }

        if owner_hash_full:
            payload["owner_hash"] = owner_hash_full

        # Attach batching_results only to the first submission (lowest context_length)
        if context_length == uploadable_single_results[0]["pp"] and batching_results:
            payload["batching_results"] = batching_results

        try:
            resp = await asyncio.to_thread(
                requests.post,
                OMLX_AI_API_URL,
                json=payload,
                timeout=15,
            )

            if resp.status_code == 201:
                data = resp.json()
                success_count += 1
                result_dict = {
                    "context_length": context_length,
                    "id": data.get("id"),
                    "url": data.get("url"),
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(
                    run,
                    {
                        "type": "upload",
                        "data": result_dict,
                    },
                )
            elif resp.status_code == 409:
                data = resp.json()
                success_count += 1  # Duplicate is still ok
                result_dict = {
                    "context_length": context_length,
                    "id": data.get("existing_id"),
                    "url": data.get("existing_url"),
                    "duplicate": True,
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(
                    run,
                    {
                        "type": "upload",
                        "data": result_dict,
                    },
                )
            else:
                failed_count += 1
                error_msg = _sanitize_upload_error(resp)
                result_dict = {
                    "context_length": context_length,
                    "error": error_msg,
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(
                    run,
                    {
                        "type": "upload",
                        "data": result_dict,
                    },
                )
                # Surface the sanitized message to ops; the full body
                # (truncated) goes to debug so it can still be retrieved
                # from the log file if needed.
                logger.warning(
                    f"Benchmark upload failed for pp{context_length}: "
                    f"{resp.status_code} {error_msg}"
                )
                if (resp.text or "")[:1] not in ("{", "["):
                    logger.debug(
                        "Benchmark upload non-JSON body (truncated): %r",
                        (resp.text or "")[:500],
                    )

        except Exception as e:
            failed_count += 1
            result_dict = {
                "context_length": context_length,
                "error": str(e),
            }
            run.upload_state["results"].append(result_dict)
            await _send_event(
                run,
                {
                    "type": "upload",
                    "data": result_dict,
                },
            )
            logger.warning(f"Benchmark upload error for pp{context_length}: {e}")

    run.upload_state["phase"] = "done"
    run.upload_state["total"] = len(uploadable_single_results)
    run.upload_state["success_count"] = success_count
    run.upload_state["failed_count"] = failed_count
    run.upload_state["skipped_count"] = skipped_count
    run.upload_state["owner_hash"] = owner_hash_display
    await _send_event(
        run,
        {
            "type": "upload_done",
            "data": {
                "owner_hash": owner_hash_display,
                "total": len(uploadable_single_results),
                "success": success_count,
                "failed": failed_count,
                "skipped": skipped_count,
                # Also on the event so SSE-only consumers (the HTML dashboard) get
                # the flags without polling /results.
                "feature_flags": run.feature_flags,
            },
        },
    )

    logger.info(
        f"Benchmark upload complete: {success_count}/"
        f"{len(uploadable_single_results)} succeeded, skipped={skipped_count}"
    )


async def run_benchmark(run: BenchmarkRun, engine_pool: Any) -> None:
    """Execute a complete benchmark run.

    Phases:
    1. Unload all loaded models
    2. Load the target model
    3. Run single request tests
    4. Run batch tests
    5. Unload the benchmark model
    """
    request = run.request
    if request.external is not None:
        await _run_external_benchmark(run)
        return
    total_tests = len(request.prompt_lengths) + len(request.batch_sizes)
    current_test = 0
    overall_start = time.perf_counter()

    # Throughput measurements must not be skewed by the memory-guard
    # throttle shrinking chunks; pin speed priority for the run.
    previous_speed_priority = _pin_speed_priority(engine_pool)

    try:
        run.model_settings_snapshot = _with_benchmark_context(
            request.context_profile,
            None,
        )
        # Snapshot experimental flags at run start. Settings can change mid-run,
        # and the produced numbers are tied to whatever was active when
        # generation actually ran.
        model_settings = None
        sm = getattr(engine_pool, "_settings_manager", None)
        if sm is not None:
            try:
                model_settings = sm.get_settings(request.model_id)
                run.experimental_features.extend(
                    _detect_experimental_features(model_settings)
                )
                run.feature_flags = _derive_feature_flags(model_settings)
                run.model_settings_snapshot = _with_benchmark_context(
                    request.context_profile,
                    _filter_uploaded_settings(model_settings),
                )
            except Exception as e:
                logger.warning(
                    f"Benchmark: failed to read experimental flags for "
                    f"{request.model_id}: {e}"
                )

        # Phase 1: Unload all loaded models
        loaded_ids = engine_pool.get_loaded_model_ids()
        if loaded_ids:
            await _send_event(
                run,
                {
                    "type": "progress",
                    "phase": "unload",
                    "message": f"Unloading {len(loaded_ids)} model(s)...",
                    "current": 0,
                    "total": total_tests,
                },
            )
            for model_id in loaded_ids:
                try:
                    await engine_pool._unload_engine(model_id)
                    logger.info(f"Benchmark: unloaded {model_id}")
                except Exception as e:
                    logger.warning(f"Benchmark: failed to unload {model_id}: {e}")

        # Phase 2: Load the target model
        await _send_event(
            run,
            {
                "type": "progress",
                "phase": "load",
                "message": f"Loading {request.model_id}...",
                "current": 0,
                "total": total_tests,
            },
        )
        # VLM MTP requires VLMBatchedEngine (which has set_vlm_mtp_drafter),
        # so don't force LM-only loading when VLM MTP is enabled.
        vlm_mtp_active = (
            model_settings is not None
            and getattr(model_settings, "vlm_mtp_enabled", False)
            and getattr(model_settings, "vlm_mtp_draft_model", None)
        )
        force_lm = True if request.force_lm_engine else not vlm_mtp_active
        engine = await engine_pool.get_engine(
            request.model_id,
            force_lm=force_lm,
        )
        logger.info(f"Benchmark: loaded {request.model_id}")

        # Generate prompts for all needed lengths
        tokenizer = engine.tokenizer
        prompts: dict[int, list[int]] = {}
        for pp_len in request.prompt_lengths:
            prompts[pp_len] = _generate_prompt(
                tokenizer,
                pp_len,
                request.context_profile,
            )

        # Ensure pp1024 prompt exists for batch tests
        if request.batch_sizes and 1024 not in prompts:
            prompts[1024] = _generate_prompt(
                tokenizer,
                1024,
                request.context_profile,
            )

        # Warmup: run a short request to trigger JIT compilation,
        # Metal shader compilation, and KV cache initialization.
        # Without this, the first real benchmark test absorbs all
        # one-time overhead and shows artificially low pp TPS.
        await _send_event(
            run,
            {
                "type": "progress",
                "phase": "warmup",
                "message": "Warming up (JIT compile)...",
                "current": 0,
                "total": total_tests,
            },
        )
        warmup_prompt = _generate_prompt(tokenizer, 32, request.context_profile)
        warmup_max_tokens = (
            request.generation_length
            if getattr(engine, "is_diffusion_model", False)
            else 8
        )
        async for _ in engine.stream_generate(
            prompt=warmup_prompt, max_tokens=warmup_max_tokens, temperature=0.0
        ):
            pass
        logger.info("Benchmark: warmup complete")

        # Start host sampling after warmup: Metal shader and JIT compilation
        # would otherwise be folded into the CPU aggregates.
        run.lifetime_footprint_at_start = get_lifetime_max_phys_footprint()
        try:
            run.sampler = SystemSampler()
            run.sampler.start()
        except Exception as e:  # noqa: BLE001
            logger.warning(f"Benchmark: host sampling unavailable: {e}")
            run.sampler = None

        # Phase 3: Single request tests
        single_pp1024_gen_tps = None

        for pp_len in request.prompt_lengths:
            current_test += 1
            await _send_event(
                run,
                {
                    "type": "progress",
                    "phase": "single",
                    "message": f"Single: pp{pp_len}/tg{request.generation_length}",
                    "current": current_test,
                    "total": total_tests,
                },
            )

            # time.monotonic only — the test internals use perf_counter, and
            # the two clocks have different epochs.
            window_start = time.monotonic()
            metrics = await _run_single_test(
                engine=engine,
                prompt=prompts[pp_len],
                max_tokens=request.generation_length,
                pp_len=pp_len,
            )
            metrics["system_metrics"] = _sample_window(run, window_start)

            result = {
                "test_type": "single",
                "pp": pp_len,
                "tg": request.generation_length,
                **metrics,
            }
            run.results.append(result)

            await _send_event(run, {"type": "result", "data": result})

            # Store pp1024 gen_tps for speedup calculation
            if pp_len == 1024:
                single_pp1024_gen_tps = metrics["gen_tps"]

        # Phase 4: Batch tests
        # Each request has a unique UUID prefix (no cache hits)
        max_batch = max(request.batch_sizes) if request.batch_sizes else 0
        batch_prompts = [
            _generate_prompt(tokenizer, 1024, request.context_profile)
            for _ in range(max_batch)
        ]

        # Skip batch tests for engines without scheduler core (e.g. VLM/Diffusion)
        batch_core = _get_batch_benchmark_core(engine)
        if request.batch_sizes and batch_core is None:
            logger.info(
                "Batch test skipped: engine does not support concurrent batching"
            )
            current_test += len(request.batch_sizes)

        for batch_size in request.batch_sizes if batch_core is not None else []:
            current_test += 1
            await _send_event(
                run,
                {
                    "type": "progress",
                    "phase": "batch",
                    "message": f"Batch {batch_size}x: pp1024/tg{request.generation_length}",
                    "current": current_test,
                    "total": total_tests,
                },
            )

            window_start = time.monotonic()
            batch_metrics = await _run_batch_test(
                engine=engine,
                prompts=batch_prompts[:batch_size],
                prompt_tokens=1024,
                max_tokens=request.generation_length,
                batch_size=batch_size,
            )
            batch_metrics["system_metrics"] = _sample_window(run, window_start)

            result = {
                "test_type": "batch",
                "pp": 1024,
                "tg": request.generation_length,
                **batch_metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Phase 5: Unload benchmark model
        await _send_event(
            run,
            {
                "type": "progress",
                "phase": "cleanup",
                "message": f"Unloading {request.model_id}...",
                "current": total_tests,
                "total": total_tests,
            },
        )
        try:
            await engine_pool._unload_engine(request.model_id)
            logger.info(f"Benchmark: unloaded {request.model_id} after benchmark")
        except Exception as e:
            logger.warning(f"Benchmark: failed to unload {request.model_id}: {e}")

        # Done
        overall_duration = time.perf_counter() - overall_start
        run.status = "completed"
        await _send_event(
            run,
            {
                "type": "done",
                "summary": {
                    "model_id": request.model_id,
                    "context_profile": request.context_profile.value,
                    "total_time": round(overall_duration, 1),
                    "total_tests": total_tests,
                },
            },
        )

        # Upload results to omlx.ai (failures don't affect benchmark status)
        try:
            await _upload_to_omlx_ai(run, engine_pool)
        except Exception as e:
            logger.warning(f"Benchmark upload to omlx.ai failed: {e}")
            await _send_event(
                run,
                {
                    "type": "upload_done",
                    "data": {
                        "owner_hash": None,
                        "total": 0,
                        "success": 0,
                        "failed": 0,
                        "error": str(e),
                    },
                },
            )

    except asyncio.CancelledError:
        run.status = "cancelled"
        await _send_event(
            run,
            {
                "type": "error",
                "message": "Benchmark cancelled by user",
            },
        )
        # Try to unload the model on cancellation
        try:
            await engine_pool._unload_engine(request.model_id)
        except Exception:
            pass

    except Exception as e:
        logger.error(f"Benchmark error: {e}", exc_info=True)
        run.status = "error"
        run.error_message = str(e)
        await _send_event(
            run,
            {
                "type": "error",
                "message": str(e),
            },
        )
        # Try to unload the model on error
        try:
            await engine_pool._unload_engine(request.model_id)
        except Exception:
            pass

    finally:
        if run.sampler is not None:
            try:
                run.sampler.stop()
            except Exception as e:  # noqa: BLE001
                logger.debug(f"Benchmark: sampler stop failed: {e}")
        _restore_speed_priority(engine_pool, previous_speed_priority)


async def _run_external_benchmark(run: BenchmarkRun) -> None:
    """Execute a benchmark run against an external OpenAI-compatible endpoint.

    No local model phases (unload/load/JIT warmup) and no community
    upload — external numbers measure someone else's hardware.
    """
    request = run.request
    total_tests = len(request.prompt_lengths) + len(request.batch_sizes)
    current_test = 0
    overall_start = time.perf_counter()
    client = ExternalAPIClient(request.external)

    try:
        # Warmup doubles as preflight: fail fast on bad URL/key and on
        # endpoints that do not return streamed usage (hard requirement
        # for accurate token counts) before any long test runs.
        await _send_event(
            run,
            {
                "type": "progress",
                "phase": "warmup",
                "message": "Warming up external endpoint...",
                "current": 0,
                "total": total_tests,
            },
        )
        await client.stream_chat_completion(
            messages=[
                {
                    "role": "user",
                    "content": _generate_external_prompt(32, request.context_profile),
                }
            ],
            max_tokens=8,
            temperature=0.0,
        )
        logger.info("Benchmark: external endpoint warmup complete")

        # Single request tests
        for pp_len in request.prompt_lengths:
            current_test += 1
            await _send_event(
                run,
                {
                    "type": "progress",
                    "phase": "single",
                    "message": f"Single: pp{pp_len}/tg{request.generation_length}",
                    "current": current_test,
                    "total": total_tests,
                },
            )

            metrics = await _run_external_single_test(
                client=client,
                prompt=_generate_external_prompt(pp_len, request.context_profile),
                max_tokens=request.generation_length,
            )

            result = {
                "test_type": "single",
                "pp": metrics["prompt_tokens"],
                "requested_pp": pp_len,
                "tg": request.generation_length,
                **metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Batch tests: concurrent requests with unique pp1024 prompts
        for batch_size in request.batch_sizes:
            current_test += 1
            await _send_event(
                run,
                {
                    "type": "progress",
                    "phase": "batch",
                    "message": f"Batch {batch_size}x: pp1024/tg{request.generation_length}",
                    "current": current_test,
                    "total": total_tests,
                },
            )

            batch_metrics = await _run_external_batch_test(
                client=client,
                prompts=[
                    _generate_external_prompt(1024, request.context_profile)
                    for _ in range(batch_size)
                ],
                max_tokens=request.generation_length,
                batch_size=batch_size,
            )

            result = {
                "test_type": "batch",
                "pp": batch_metrics["prompt_tokens"],
                "requested_pp": 1024,
                "tg": request.generation_length,
                **batch_metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Done
        overall_duration = time.perf_counter() - overall_start
        run.status = "completed"
        await _send_event(
            run,
            {
                "type": "done",
                "summary": {
                    "model_id": request.model_id,
                    "context_profile": request.context_profile.value,
                    "total_time": round(overall_duration, 1),
                    "total_tests": total_tests,
                },
            },
        )

        # External results measure remote hardware — never upload them to
        # the omlx.ai community leaderboard. Mirrors the experimental-
        # features skip so REST pollers see the same upload_state shape.
        run.upload_state["phase"] = "skipped"
        run.upload_state["skipped_reason"] = "external_endpoint"
        await _send_event(
            run,
            {
                "type": "upload_skipped",
                "reason": "external_endpoint",
                "features": [],
            },
        )

    except asyncio.CancelledError:
        run.status = "cancelled"
        await _send_event(
            run,
            {
                "type": "error",
                "message": "Benchmark cancelled by user",
            },
        )
    except Exception as e:
        logger.error(f"External benchmark error: {e}", exc_info=True)
        run.status = "error"
        run.error_message = str(e)
        await _send_event(
            run,
            {
                "type": "error",
                "message": str(e),
            },
        )
    finally:
        await client.aclose()


# ===== FILE: omlx/admin/build_css.py =====

#!/usr/bin/env python3
"""Build script for Tailwind CSS compilation.

Downloads Tailwind v3 standalone CLI if needed and compiles CSS.
Requires no Node.js installation.

Usage:
    cd omlx/omlx/admin
    python build_css.py          # Build minified CSS
    python build_css.py --watch  # Watch mode for development
"""

import platform
import subprocess
import sys
import urllib.request
from pathlib import Path

TAILWIND_VERSION = "v3.4.17"
ADMIN_DIR = Path(__file__).parent


def get_binary_name() -> str:
    """Get platform-specific Tailwind CLI binary name."""
    machine = platform.machine().lower()
    system = platform.system().lower()
    if system == "darwin":
        arch = "arm64" if machine == "arm64" else "x64"
        return f"tailwindcss-macos-{arch}"
    elif system == "linux":
        arch = "arm64" if "aarch64" in machine else "x64"
        return f"tailwindcss-linux-{arch}"
    raise RuntimeError(f"Unsupported platform: {system} {machine}")


def ensure_binary() -> Path:
    """Download Tailwind standalone CLI if not present."""
    binary_name = get_binary_name()
    binary_path = ADMIN_DIR / binary_name

    if binary_path.exists():
        return binary_path

    url = (
        f"https://github.com/tailwindlabs/tailwindcss/releases/download/"
        f"{TAILWIND_VERSION}/{binary_name}"
    )
    print(f"Downloading Tailwind CSS {TAILWIND_VERSION}...")
    print(f"  {url}")
    urllib.request.urlretrieve(url, binary_path)
    binary_path.chmod(0o755)
    print(f"  Saved to {binary_path}")
    return binary_path


def main() -> None:
    binary = ensure_binary()

    input_css = ADMIN_DIR / "src" / "input.css"
    output_css = ADMIN_DIR / "static" / "css" / "tailwind.css"
    config = ADMIN_DIR / "tailwind.config.js"

    output_css.parent.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(binary),
        "-i", str(input_css),
        "-o", str(output_css),
        "-c", str(config),
        "--minify",
    ]

    if "--watch" in sys.argv:
        cmd.append("--watch")
        print("Watching for changes... (Ctrl+C to stop)")

    print(f"Building: {' '.join(cmd)}")
    result = subprocess.run(cmd)

    if result.returncode == 0 and "--watch" not in sys.argv:
        size = output_css.stat().st_size
        print(f"Output: {output_css} ({size:,} bytes)")
    sys.exit(result.returncode)


if __name__ == "__main__":
    main()


# ===== FILE: omlx/admin/context_benchmark.py =====

# SPDX-License-Identifier: Apache-2.0
"""Context window benchmark for oMLX admin panel.

Measures the largest prompt the current machine can actually prefill for
a model, then writes the result into the model's ``max_context_window``
setting so clients get an honest upfront limit instead of mid-prefill
aborts.

Method: a short calibration prefill seeds the scheduler's transient
tracker, then the deterministic admission boundary is found by bisecting
``scheduler.preflight_or_raise`` in-process — microseconds per probe, no
GPU work. One real prefill at the (2k-floored) boundary verifies it end
to end; on a mid-prefill abort a single conservative retry runs, sized
from where the first prefill actually died (90% of its processed
tokens). A clean first-try completion instead climbs upward in 1.2x
extension probes until the cap or the first failure, since a boundary
that binds near the ceiling can be conservative. Probe requests carry
``skip_cache_store`` and the bench clears the model's paged cache
between probes, so nothing leaks into the prefix/SSD cache tiers.
"""

import asyncio
import contextlib
import gc
import logging
import time
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from pydantic import BaseModel, field_validator

from ..exceptions import PrefillMemoryAbortedError, PrefillMemoryExceededError
from ..prefill_progress import get_prefill_tracker
from .benchmark import _generate_prompt

logger = logging.getLogger(__name__)

# Module-level storage for context benchmark runs
_context_runs: dict[str, "ContextBenchmarkRun"] = {}

# Valid measurement targets (16k .. 512k)
VALID_TARGET_TOKENS = [16384, 32768, 65536, 131072, 262144, 524288]

# Applied values are floored to this granularity (user-facing "2k units").
_APPLY_GRANULARITY = 2048

# Smallest boundary worth reporting; below this the machine cannot hold a
# usable context for the model at all.
_MIN_USEFUL_TOKENS = 2048

# Calibration prefill size — seeds the transient tracker EWMA and the
# GDN/Mamba fixed-state measurement before the analytic search.
_CALIBRATION_TOKENS = 4096

# Two real-prefill attempts total: the boundary try, then one
# conservative retry informed by where the first prefill actually
# aborted. More rounds add minutes of near-ceiling crawling for
# marginal precision. Instant admission rejections (nothing prefilled)
# do not consume an attempt; the spin cap bounds those instead.
_MAX_VERIFY_ATTEMPTS = 2
_MAX_VERIFY_SPINS = 6

# Retry candidate = this fraction of the token count the aborted prefill
# actually completed — physical evidence of what fits, with headroom.
_ABORT_EVIDENCE_SAFETY = 0.9

# When the verify completed straight at the analytic boundary (no abort
# evidence, memory-bound), extension probes climb in steps of this
# factor until the cap or the first failure — a conservative boundary
# gets the chance to prove more. Only COMPLETED extensions are ever
# applied; a failed one keeps the last completed value.
_EXTENSION_FACTOR = 1.2

_CTX_TERMINAL_TYPES = frozenset({"done", "error"})

# Overall progress bands per phase (0-100).
_PROGRESS_BANDS = {
    "prepare": (0.0, 8.0),
    "calibrate": (8.0, 16.0),
    "estimate": (16.0, 20.0),
    "verify": (20.0, 95.0),
    "apply": (95.0, 100.0),
}


class ContextBenchmarkRequest(BaseModel):
    """Request model for starting a context benchmark."""

    model_id: str
    target_tokens: int = 131072

    @field_validator("target_tokens")
    @classmethod
    def validate_target_tokens(cls, v: int) -> int:
        if v not in VALID_TARGET_TOKENS:
            raise ValueError(
                f"Invalid target {v}. Must be one of {VALID_TARGET_TOKENS}"
            )
        return v


@dataclass
class ContextBenchmarkRun:
    """Tracks the state of a running context benchmark.

    Same SSE delivery model as ``BenchmarkRun``: events are appended to
    `events` under `cond`, subscribers replay from offset 0 then wait for
    new entries, and `terminal` closes the stream. `phase` / `progress` /
    `message` mirror the latest progress event so REST pollers (the
    native app) never need to parse the event log.
    """

    bench_id: str
    request: ContextBenchmarkRequest
    status: str = "running"  # running, completed, cancelled, error
    events: list[dict] = field(default_factory=list)
    cond: asyncio.Condition = field(default_factory=asyncio.Condition)
    terminal: bool = False
    task: asyncio.Task | None = None
    phase: str = "prepare"
    progress: float = 0.0
    message: str = ""
    result: dict | None = None
    error_message: str = ""


def get_run(bench_id: str) -> "ContextBenchmarkRun | None":
    """Get a context benchmark run by ID."""
    return _context_runs.get(bench_id)


def get_active_run() -> "ContextBenchmarkRun | None":
    """Return the currently-running context benchmark, if any."""
    for run in _context_runs.values():
        if run.status == "running":
            return run
    return None


def create_run(request: ContextBenchmarkRequest) -> ContextBenchmarkRun:
    """Create and register a new context benchmark run."""
    bench_id = f"ctx-{uuid.uuid4().hex[:12]}"
    run = ContextBenchmarkRun(bench_id=bench_id, request=request)
    _context_runs[bench_id] = run
    return run


def cleanup_old_runs(max_runs: int = 10) -> None:
    """Remove old completed runs to prevent memory leaks."""
    completed = [
        (bid, r)
        for bid, r in _context_runs.items()
        if r.status in ("completed", "cancelled", "error")
    ]
    if len(completed) > max_runs:
        for bid, _ in completed[:-max_runs]:
            del _context_runs[bid]


async def _send_event(run: ContextBenchmarkRun, event: dict) -> None:
    """Append an event, mirror progress state, wake subscribers."""
    async with run.cond:
        run.events.append(event)
        if event.get("type") == "progress":
            run.phase = event.get("phase", run.phase)
            run.progress = float(event.get("progress", run.progress))
            run.message = event.get("message", run.message)
        if event.get("type") in _CTX_TERMINAL_TYPES:
            run.terminal = True
        run.cond.notify_all()


async def _progress(
    run: ContextBenchmarkRun, phase: str, fraction: float, message: str
) -> None:
    """Emit a progress event mapped into the phase's overall band."""
    lo, hi = _PROGRESS_BANDS.get(phase, (0.0, 100.0))
    overall = lo + (hi - lo) * min(1.0, max(0.0, fraction))
    await _send_event(
        run,
        {
            "type": "progress",
            "phase": phase,
            "progress": round(overall, 1),
            "message": message,
        },
    )


def floor_to_apply_granularity(tokens: int) -> int:
    """Floor a token count to the applied 2k granularity."""
    return (max(0, tokens) // _APPLY_GRANULARITY) * _APPLY_GRANULARITY


def ceil_to_apply_granularity(tokens: int) -> int:
    """Ceil a token count to the applied 2k granularity."""
    return -(-max(0, tokens) // _APPLY_GRANULARITY) * _APPLY_GRANULARITY


def bisect_admission(fits: Callable[[int], bool], lo: int, hi: int) -> int:
    """Largest n in [lo, hi] with fits(n) True; 0 if even lo fails.

    Assumes fits is monotone non-increasing over n (the admission
    formula is: estimate grows with n, limit is fixed).
    """
    if hi < lo:
        return 0
    if not fits(lo):
        return 0
    if fits(hi):
        return hi
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if fits(mid):
            lo = mid
        else:
            hi = mid
    return lo


def _resolve_scheduler(engine: Any) -> Any | None:
    """Pool engine → async engine → core → scheduler (None when absent)."""
    async_core = getattr(engine, "_engine", None)
    core = getattr(async_core, "engine", None) if async_core is not None else None
    return getattr(core, "scheduler", None) if core is not None else None


def _resolve_core(engine: Any) -> Any | None:
    async_core = getattr(engine, "_engine", None)
    return getattr(async_core, "engine", None) if async_core is not None else None


def _guard_ready(scheduler: Any) -> bool:
    """preflight_or_raise no-ops unless all three of these hold."""
    return bool(
        getattr(scheduler, "_prefill_memory_guard", False)
        and getattr(scheduler, "_memory_hard_limit_bytes", 0) > 0
        and getattr(scheduler, "memory_monitor", None) is not None
    )


def _make_fits(scheduler: Any) -> Callable[[int], bool]:
    def fits(n: int) -> bool:
        try:
            scheduler.preflight_or_raise(num_prompt_tokens=n)
        except PrefillMemoryExceededError:
            return False
        return True

    return fits


def _clear_probe_residue(scheduler: Any) -> None:
    """Release probe KV blocks and pooled Metal buffers.

    Runs ON the engine's MLX executor thread (thread-local stream +
    serialization with scheduler steps). The bench loads its model fresh
    with everything else unloaded, so clearing the paged cache only drops
    probe blocks — there is no user prefix cache to lose.
    """
    from ..scheduler import _sync_and_clear_cache

    cache = getattr(scheduler, "block_aware_cache", None)
    if cache is not None:
        try:
            cache.clear()
        except Exception as exc:
            logger.debug("Context bench: paged cache clear failed: %s", exc)
    gc.collect()
    _sync_and_clear_cache(getattr(scheduler, "_stream", None))
    # Refresh the cached mlx-active sample so event-loop preflight probes
    # see post-cleanup usage instead of the prefill-peak snapshot.
    try:
        scheduler._current_usage_bytes()
    except Exception as exc:
        logger.debug("Context bench: usage refresh failed: %s", exc)


async def _cleanup_between_probes(engine: Any, scheduler: Any) -> None:
    """Dispatch probe-residue cleanup to the engine's executor and settle."""
    core = _resolve_core(engine)
    executor = getattr(core, "_mlx_executor", None)
    loop = asyncio.get_running_loop()
    for _ in range(2):
        if executor is not None:
            await loop.run_in_executor(executor, _clear_probe_residue, scheduler)
        await asyncio.sleep(0.5)


async def _run_probe_prefill(engine: Any, prompt: str) -> tuple[Any, float]:
    """Prefill the prompt with max_tokens=1; return (final output, seconds).

    Raises PrefillMemoryExceededError / PrefillMemoryAbortedError when the
    memory guard refuses or aborts the probe. With max_tokens=1 the wall
    clock is dominated by the prefill, so callers can derive a prefill
    tok/s from it when the engine does not report prompt_tps.
    """
    started = time.perf_counter()
    last_output = None
    async for output in engine.stream_generate(
        prompt=prompt,
        max_tokens=1,
        temperature=0.0,
        top_p=1.0,
        skip_cache_store=True,
    ):
        last_output = output
    return last_output, time.perf_counter() - started


def next_verify_candidate(
    candidate: int, observed_processed: int, new_boundary: int
) -> int:
    """Conservative retry candidate after a failed verify prefill.

    Physical evidence caps from above: the aborted prefill completed
    ``observed_processed`` tokens before dying, so 90% of that is the
    most a retry should attempt. The re-measured boundary caps too —
    the failure-path cleanup resets the transient tracker first, so the
    re-bisect reflects honest KV/static pricing at the current baseline
    rather than the dead prefill's contaminated last-chunk delta.
    Without evidence (nothing prefilled), halve.
    """
    if observed_processed > 0:
        nxt = min(
            int(observed_processed * _ABORT_EVIDENCE_SAFETY),
            candidate - _APPLY_GRANULARITY,
        )
    else:
        nxt = min(candidate // 2, candidate - _APPLY_GRANULARITY)
    if new_boundary > 0:
        nxt = min(nxt, new_boundary)
    return floor_to_apply_granularity(nxt)


def _reset_transient_tracker(scheduler: Any) -> None:
    """Drop the scheduler's chunk-transient measurements after a failed probe.

    An enforcer-killed prefill leaves a giant last-chunk phys delta in the
    tracker; under speed priority that poisons every subsequent admission
    (a 2048-token charge of several GB), and since rejected probes never
    run a chunk, the poison would never decay within the bench. Resetting
    returns admission to the static + exact-KV estimate until real chunks
    re-seed the EWMA.
    """
    tracker = getattr(scheduler, "_prefill_transient_tracker", None)
    if tracker is not None:
        try:
            tracker.reset()
        except Exception as exc:
            logger.debug("Context bench: tracker reset failed: %s", exc)


async def _relay_prefill_progress(
    run: ContextBenchmarkRun,
    model_id: str,
    attempt: int,
    candidate: int,
    progress_holder: dict,
) -> None:
    """Poll the global prefill tracker and mirror % into verify progress.

    Also records the last seen processed-token count into
    ``progress_holder`` so a failed attempt can derive its retry
    candidate from where the prefill actually died.
    """
    tracker = get_prefill_tracker()
    last_emitted = -1.0
    while True:
        await asyncio.sleep(1.5)
        entries = tracker.get_model_progress(model_id)
        if not entries:
            continue
        entry = entries[0]
        progress_holder["processed"] = max(
            int(progress_holder.get("processed", 0)), int(entry["processed"])
        )
        total = max(1, entry["total"])
        fraction = entry["processed"] / total
        if fraction - last_emitted < 0.01:
            continue
        last_emitted = fraction
        message = (
            f"Verify prefill {candidate:,} tokens (attempt {attempt}): "
            f"{entry['processed']:,}/{entry['total']:,}"
        )
        details = []
        speed = entry.get("speed")
        if speed:
            details.append(f"{speed:,.0f} tok/s")
        eta = entry.get("eta")
        if eta is not None:
            details.append(f"~{int(eta)}s left")
        if details:
            message += f" ({', '.join(details)})"
        await _progress(run, "verify", fraction, message)


async def run_context_benchmark(run: ContextBenchmarkRun, engine_pool: Any) -> None:
    """Execute a context benchmark run.

    Phases: prepare (unload all → load target) → calibrate → estimate
    (bisect admission) → verify (real prefill, step down on abort) →
    apply (write max_context_window) → cleanup (unload).
    """
    request = run.request
    overall_start = time.perf_counter()

    try:
        # Phase 1: prepare — unload everything, load the target model.
        loaded_ids = engine_pool.get_loaded_model_ids()
        if loaded_ids:
            await _progress(
                run, "prepare", 0.1, f"Unloading {len(loaded_ids)} model(s)..."
            )
            for model_id in loaded_ids:
                try:
                    await engine_pool._unload_engine(model_id)
                    logger.info("Context bench: unloaded %s", model_id)
                except Exception as exc:
                    logger.warning(
                        "Context bench: failed to unload %s: %s", model_id, exc
                    )

        await _progress(run, "prepare", 0.4, f"Loading {request.model_id}...")
        engine = await engine_pool.get_engine(request.model_id)
        logger.info("Context bench: loaded %s", request.model_id)

        scheduler = _resolve_scheduler(engine)
        if scheduler is None:
            raise RuntimeError(
                "This model's engine does not expose a scheduler, so the "
                "context benchmark cannot probe its admission boundary. "
                "Disable speculative engine features (e.g. DFlash) for this "
                "model and retry."
            )
        if not _guard_ready(scheduler):
            raise RuntimeError(
                "The prefill memory guard is disabled, so there is no "
                "admission boundary to measure. Enable Memory Guard and "
                "retry."
            )

        entry = engine_pool.get_entry(request.model_id)
        native = getattr(entry, "model_context_length", None) or 0
        cap = request.target_tokens
        cap_source = "target"
        if 0 < native < cap:
            cap = native
            cap_source = "native"

        # Phase 2: calibrate — JIT warmup, then one real mid-size prefill
        # to seed the transient tracker EWMA and the GDN fixed-state probe.
        tokenizer = engine.tokenizer
        await _progress(run, "calibrate", 0.1, "Warming up (JIT compile)...")
        async for _ in engine.stream_generate(
            prompt=_generate_prompt(tokenizer, 32),
            max_tokens=8,
            temperature=0.0,
            skip_cache_store=True,
        ):
            pass

        calibration_tokens = min(_CALIBRATION_TOKENS, max(1024, cap // 2))
        await _progress(
            run,
            "calibrate",
            0.4,
            f"Calibration prefill ({calibration_tokens:,} tokens)...",
        )
        try:
            await _run_probe_prefill(
                engine, _generate_prompt(tokenizer, calibration_tokens)
            )
        except (PrefillMemoryExceededError, PrefillMemoryAbortedError) as exc:
            raise RuntimeError(
                f"Not enough memory to prefill even {calibration_tokens:,} "
                f"tokens on this machine: {exc}"
            ) from exc
        await _cleanup_between_probes(engine, scheduler)

        # Phase 3: estimate — bisect the deterministic admission boundary.
        await _progress(run, "estimate", 0.3, "Estimating admission boundary...")
        fits = _make_fits(scheduler)
        boundary = bisect_admission(fits, 1024, cap)
        if floor_to_apply_granularity(boundary) < _MIN_USEFUL_TOKENS:
            raise RuntimeError(
                "Admission boundary is below 2k tokens — not enough free "
                "memory to serve a usable context for this model."
            )
        logger.info(
            "Context bench: admission boundary %d tokens (cap %d, %s)",
            boundary,
            cap,
            cap_source,
        )

        # Phase 4: verify — real prefill at the floored candidate. On a
        # mid-prefill abort, one conservative retry sized from where the
        # first prefill actually died (90% of its processed tokens). An
        # INSTANT admission rejection (0 tokens processed — current drifted
        # between the bisection and the probe) costs no GPU time and does
        # not consume a real-prefill attempt; it just re-bisects with the
        # fresh usage and retries, bounded by the spin cap.
        candidate = floor_to_apply_granularity(min(boundary, cap))
        verified = 0
        verified_prompt_tokens = 0
        verify_prefill_tps = 0.0
        attempts = 0
        spins = 0
        had_abort = False
        while spins < _MAX_VERIFY_SPINS:
            spins += 1
            attempt_no = attempts + 1
            await _progress(
                run,
                "verify",
                0.0,
                f"Verify prefill {candidate:,} tokens (attempt {attempt_no})...",
            )
            progress_holder: dict = {}
            relay = asyncio.create_task(
                _relay_prefill_progress(
                    run, request.model_id, attempt_no, candidate, progress_holder
                )
            )
            try:
                output, probe_seconds = await _run_probe_prefill(
                    engine, _generate_prompt(tokenizer, candidate)
                )
            except (PrefillMemoryExceededError, PrefillMemoryAbortedError) as exc:
                observed = int(progress_holder.get("processed", 0))
                instant_reject = (
                    not isinstance(exc, PrefillMemoryAbortedError) and observed == 0
                )
                if not instant_reject:
                    attempts += 1
                    had_abort = True
                logger.info(
                    "Context bench: verify failed at %d tokens "
                    "(attempt %d, %d processed%s): %s",
                    candidate,
                    attempt_no,
                    observed,
                    ", instant reject" if instant_reject else "",
                    exc,
                )
                if attempts >= _MAX_VERIFY_ATTEMPTS:
                    raise RuntimeError(
                        f"Verification failed {_MAX_VERIFY_ATTEMPTS} times — "
                        f"the machine could not complete a prefill at the "
                        f"measured sizes. Raise the Memory Guard ceiling or "
                        f"pick a smaller target and rerun."
                    ) from exc
                await _cleanup_between_probes(engine, scheduler)
                # Drop the dead prefill's contaminated transient sample so
                # the re-bisect prices the next candidate honestly.
                _reset_transient_tracker(scheduler)
                new_boundary = bisect_admission(fits, 1024, candidate)
                if instant_reject:
                    # Free retry: one grain below the freshly measured
                    # boundary.
                    candidate = floor_to_apply_granularity(
                        min(new_boundary, candidate - _APPLY_GRANULARITY)
                    )
                else:
                    candidate = next_verify_candidate(candidate, observed, new_boundary)
                if candidate < _MIN_USEFUL_TOKENS:
                    raise RuntimeError(
                        "Verification kept failing above the 2k floor — "
                        "not enough stable memory for a usable context."
                    ) from exc
                continue
            finally:
                relay.cancel()
            attempts += 1
            verified = candidate
            verified_prompt_tokens = int(
                getattr(output, "prompt_tokens", 0) or candidate
            )
            # Prefill speed of the successful verify: prefer the engine's
            # own figure, fall back to the probe wall clock (max_tokens=1,
            # so it is essentially all prefill).
            verify_prefill_tps = float(getattr(output, "prompt_tps", 0.0) or 0.0)
            if verify_prefill_tps <= 0 and probe_seconds > 0:
                verify_prefill_tps = verified_prompt_tokens / probe_seconds
            break
        if verified == 0:
            raise RuntimeError(
                "Verification did not converge within the retry budget. "
                "Raise the Memory Guard ceiling or pick a smaller target "
                "and rerun."
            )

        # Extension probes: the verified value came straight from the
        # analytic boundary (clean first-try completion, memory-bound).
        # Near the ceiling the boundary can be conservative, so keep
        # probing at 1.2x steps until the cap is reached or a probe
        # fails. Every probe re-passes the live admission gate at the
        # settled baseline, and only COMPLETED prefills raise the value;
        # the first rejection or abort ends the climb and keeps the last
        # completed size.
        extended = False
        extension_aborted = False
        while not had_abort and verified < floor_to_apply_granularity(cap):
            ext_target = min(
                ceil_to_apply_granularity(int(verified * _EXTENSION_FACTOR)), cap
            )
            if ext_target <= verified:
                break
            await _cleanup_between_probes(engine, scheduler)
            await _progress(
                run,
                "verify",
                0.98,
                f"Extension probe {ext_target:,} tokens " f"({_EXTENSION_FACTOR}x)...",
            )
            progress_holder = {}
            relay = asyncio.create_task(
                _relay_prefill_progress(
                    run,
                    request.model_id,
                    attempts + 1,
                    ext_target,
                    progress_holder,
                )
            )
            try:
                output, probe_seconds = await _run_probe_prefill(
                    engine, _generate_prompt(tokenizer, ext_target)
                )
            except (
                PrefillMemoryExceededError,
                PrefillMemoryAbortedError,
            ) as exc:
                observed = int(progress_holder.get("processed", 0))
                if observed > 0 or isinstance(exc, PrefillMemoryAbortedError):
                    attempts += 1
                    extension_aborted = True
                logger.info(
                    "Context bench: extension probe failed at %d tokens "
                    "(%d processed): %s",
                    ext_target,
                    observed,
                    exc,
                )
                await _cleanup_between_probes(engine, scheduler)
                _reset_transient_tracker(scheduler)
                break
            else:
                attempts += 1
                extended = True
                verified = ext_target
                verified_prompt_tokens = int(
                    getattr(output, "prompt_tokens", 0) or ext_target
                )
                verify_prefill_tps = float(getattr(output, "prompt_tps", 0.0) or 0.0)
                if verify_prefill_tps <= 0 and probe_seconds > 0:
                    verify_prefill_tps = verified_prompt_tokens / probe_seconds
                logger.info(
                    "Context bench: extension probe completed at %d tokens",
                    ext_target,
                )
            finally:
                relay.cancel()

        # Phase 5: apply — re-measure with the post-verify tracker state
        # (the near-ceiling prefill can raise the floor-chunk transient
        # charge) and write the tighter of the two, floored to 2k.
        await _progress(run, "apply", 0.2, "Re-checking admission boundary...")
        await _cleanup_between_probes(engine, scheduler)
        post_boundary = bisect_admission(fits, 1024, verified)
        final = verified
        if post_boundary > 0:
            # The re-bisect can be contaminated by the verify prefill's own
            # residue (buffer pool, last-chunk transient sample) — never
            # tighten below 90% of what physically completed. After a
            # FAILED extension probe the contamination is fresh and the
            # verified size completed moments ago, so it stands as-is.
            if extension_aborted:
                evidence_floor = floor_to_apply_granularity(verified)
            else:
                evidence_floor = floor_to_apply_granularity(
                    int(verified * _ABORT_EVIDENCE_SAFETY)
                )
            final = min(
                final,
                max(floor_to_apply_granularity(post_boundary), evidence_floor),
            )
        final = floor_to_apply_granularity(final)
        if final < _MIN_USEFUL_TOKENS:
            final = _MIN_USEFUL_TOKENS

        capped_by = "memory"
        if final >= floor_to_apply_granularity(cap):
            capped_by = cap_source

        applied = False
        sm = getattr(engine_pool, "_settings_manager", None)
        if sm is not None:
            try:
                settings = sm.get_settings(request.model_id)
                settings.max_context_window = final
                sm.set_settings(request.model_id, settings)
                applied = True
                logger.info(
                    "Context bench: applied max_context_window=%d for %s",
                    final,
                    request.model_id,
                )
            except Exception as exc:
                logger.warning(
                    "Context bench: failed to apply setting for %s: %s",
                    request.model_id,
                    exc,
                )
        await _progress(run, "apply", 0.8, "Applying context window setting...")

        result = {
            "model_id": request.model_id,
            "target_tokens": request.target_tokens,
            "native_context_length": native or None,
            "measured_tokens": boundary,
            "verified_tokens": verified,
            "verified_prompt_tokens": verified_prompt_tokens,
            "applied_tokens": final,
            "applied": applied,
            "capped_by": capped_by,
            "attempts": attempts,
            # True when the 1.2x extension probe completed and raised the
            # verified value beyond the analytic boundary.
            "extended": extended,
            # Prefill tok/s of the successful verify run — what a prompt at
            # the applied size actually prefills at in the current mode.
            "prefill_tps": round(verify_prefill_tps, 1),
            "duration_s": round(time.perf_counter() - overall_start, 1),
            # Mode the measurement ran under — the applied value only holds
            # while serving keeps the same prefill priority.
            "prefill_priority": (
                "speed"
                if getattr(scheduler, "_prefill_speed_priority", False)
                else "context"
            ),
        }
        run.result = result
        await _send_event(run, {"type": "result", "data": result})

        # Phase 6: cleanup — unload the bench model (throughput parity).
        try:
            await engine_pool._unload_engine(request.model_id)
            logger.info("Context bench: unloaded %s after run", request.model_id)
        except Exception as exc:
            logger.warning(
                "Context bench: failed to unload %s: %s", request.model_id, exc
            )

        run.status = "completed"
        await _send_event(
            run,
            {
                "type": "done",
                "summary": {
                    "model_id": request.model_id,
                    "applied_tokens": final,
                    "total_time": result["duration_s"],
                },
            },
        )

    except asyncio.CancelledError:
        run.status = "cancelled"
        run.error_message = "Context benchmark cancelled by user"
        await _send_event(run, {"type": "error", "message": run.error_message})
        with contextlib.suppress(Exception):
            await engine_pool._unload_engine(request.model_id)

    except Exception as exc:
        logger.error("Context bench error: %s", exc, exc_info=True)
        run.status = "error"
        run.error_message = str(exc)
        await _send_event(run, {"type": "error", "message": str(exc)})
        with contextlib.suppress(Exception):
            await engine_pool._unload_engine(request.model_id)


# ===== FILE: omlx/admin/external_api.py =====

# SPDX-License-Identifier: Apache-2.0
"""OpenAI-compatible external endpoint client for admin benchmarks.

Shared by the throughput and accuracy benchmarks to run against a remote
/chat/completions endpoint instead of a local engine. Token counts always
come from the endpoint's usage payload — SSE chunks are never counted as
tokens because providers batch multiple tokens per chunk.
"""

import json
import logging
import re
import time
from dataclasses import dataclass
from typing import Any, Optional

import httpx
from pydantic import BaseModel, Field, SecretStr, field_validator

logger = logging.getLogger(__name__)

# read=3600 covers both the largest between-chunk gap on streams (TTFT of a
# very long prefill on a slow remote) and the full-response wait for
# non-streaming accuracy calls. Benchmarks are supervised and cancellable,
# so a generous ceiling beats spurious failures; connect=15 still fails
# dead endpoints fast.
DEFAULT_TIMEOUT = httpx.Timeout(connect=15.0, read=3600.0, write=120.0, pool=30.0)

_ERROR_DETAIL_MAX_CHARS = 300

_REASONING_FIELD_NAMES = ("reasoning_content", "reasoning", "analysis")

# Thinking models spend reasoning tokens before message.content, and those
# tokens count toward max_tokens on OpenAI-compatible APIs, so a small cap
# truncates them before they can answer (#2309). max_tokens only bounds
# runaway generation; non-thinking endpoints still stop after a few tokens.
_PREFLIGHT_MAX_TOKENS = 4096

# Provider-specific request JSON must not override fields owned by the
# benchmark or authentication layer.
PROTECTED_EXTRA_BODY_FIELDS = frozenset({
    "model",
    "messages",
    "stream",
    "stream_options",
    "max_tokens",
    "temperature",
    "api_key",
    "authorization",
})


class ExternalEndpointConfig(BaseModel):
    """Connection settings for an external OpenAI-compatible endpoint.

    api_key is a SecretStr so the key never leaks through repr() or logs.
    """

    base_url: str
    api_key: SecretStr = SecretStr("")
    model: str
    extra_body: dict[str, Any] = Field(default_factory=dict)

    @field_validator("base_url")
    @classmethod
    def validate_base_url(cls, v: str) -> str:
        v = v.strip().rstrip("/")
        if not v.startswith(("http://", "https://")):
            raise ValueError("base_url must start with http:// or https://")
        return v

    @field_validator("model")
    @classmethod
    def validate_model(cls, v: str) -> str:
        v = v.strip()
        if not v:
            raise ValueError("model must not be empty")
        return v

    @field_validator("extra_body", mode="before")
    @classmethod
    def validate_extra_body(cls, v: Any) -> dict[str, Any]:
        if v is None:
            return {}
        if not isinstance(v, dict):
            raise ValueError("extra_body must be a JSON object")
        blocked = sorted(
            str(key)
            for key in v
            if str(key).lower() in PROTECTED_EXTRA_BODY_FIELDS
        )
        if blocked:
            raise ValueError(
                "extra_body cannot override protected field(s): "
                + ", ".join(blocked)
            )
        return dict(v)


class ExternalEndpointError(Exception):
    """User-presentable failure talking to an external endpoint."""

    def __init__(self, message: str, status: str = "invalid_response"):
        super().__init__(message)
        self.status = status


@dataclass
class StreamStats:
    """Timing and token stats from one streamed chat completion."""

    prompt_tokens: int
    completion_tokens: int
    cached_tokens: int
    start_time: float
    first_content_time: float
    last_content_time: float
    end_time: float
    text: str
    # False when the stream carried no content or reasoning delta at all.
    # first/last_content_time then fall back to end_time, so every timing
    # derived from them (TTFT, prefill rate, decode rate) describes the
    # whole response instead of the phase it claims to measure.
    content_observed: bool = True


@dataclass
class ChatResult:
    """Non-streaming chat completion result."""

    text: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
    finish_reason: Optional[str] = None
    status: str = "ok"
    reasoning_fields_present: tuple[str, ...] = ()
    reasoning_fields_nonempty: tuple[str, ...] = ()


def _extract_error_detail(body: str) -> str:
    """Pull a short human-readable message out of an error response body."""
    try:
        data = json.loads(body)
        if isinstance(data, dict):
            err = data.get("error")
            if isinstance(err, dict) and err.get("message"):
                return str(err["message"])[:_ERROR_DETAIL_MAX_CHARS]
            for key in ("message", "detail"):
                if data.get(key):
                    return str(data[key])[:_ERROR_DETAIL_MAX_CHARS]
    except ValueError:
        pass
    text = body.strip()
    if "<" in text and ">" in text:
        return f"unexpected non-JSON response ({len(body)} bytes)"
    return text[:_ERROR_DETAIL_MAX_CHARS] or "no response body"


class ExternalAPIClient:
    """Async client for an external OpenAI-compatible /chat/completions API.

    Provider-specific fields may be supplied through config.extra_body. The
    config validator prevents them from overriding benchmark-owned or
    authentication-related fields.
    """

    def __init__(
        self,
        config: ExternalEndpointConfig,
        timeout: httpx.Timeout = DEFAULT_TIMEOUT,
        transport: Optional[httpx.AsyncBaseTransport] = None,
    ):
        self._config = config
        self._chat_url = f"{config.base_url}/chat/completions"
        headers = {}
        key = config.api_key.get_secret_value()
        if key:
            headers["Authorization"] = f"Bearer {key}"
        # transport is injectable for tests (httpx.MockTransport).
        self._client = httpx.AsyncClient(
            headers=headers,
            timeout=timeout,
            limits=httpx.Limits(max_connections=64),
            transport=transport,
        )

    async def aclose(self) -> None:
        await self._client.aclose()

    def _build_body(
        self,
        messages: list[dict],
        max_tokens: int,
        temperature: Optional[float],
        stream: bool,
    ) -> dict:
        body: dict[str, Any] = {
            "model": self._config.model,
            "messages": messages,
            "max_tokens": max_tokens,
        }
        if temperature is not None:
            body["temperature"] = temperature
        if stream:
            body["stream"] = True
            body["stream_options"] = {"include_usage": True}
        body.update(self._config.extra_body)
        return body

    def _redact_secret(self, text: str) -> str:
        """Defensively remove the configured key from provider errors."""
        key = self._config.api_key.get_secret_value()
        if key:
            return text.replace(key, "[REDACTED]")
        return text

    def _map_transport_error(self, exc: httpx.HTTPError) -> ExternalEndpointError:
        base_url = self._config.base_url
        if isinstance(exc, httpx.TimeoutException):
            if isinstance(exc, httpx.ConnectTimeout):
                return ExternalEndpointError(
                    f"Timed out connecting to external endpoint {base_url}",
                    status="timeout",
                )
            return ExternalEndpointError(
                "External endpoint timed out while waiting for a response",
                status="timeout",
            )
        if isinstance(exc, httpx.ConnectError):
            return ExternalEndpointError(
                self._redact_secret(
                    f"Cannot connect to external endpoint {base_url}: {exc}"
                ),
                status="connection_error",
            )
        return ExternalEndpointError(
            self._redact_secret(
                f"External endpoint request failed: {type(exc).__name__}: {exc}"
            ),
            status="connection_error",
        )

    def _status_error(self, status: int, body_text: str) -> ExternalEndpointError:
        if status in (401, 403):
            return ExternalEndpointError(
                f"External endpoint rejected the API key (HTTP {status})",
                status="http_error",
            )
        detail = self._redact_secret(_extract_error_detail(body_text))
        return ExternalEndpointError(
            f"External endpoint returned HTTP {status}: {detail}",
            status="http_error",
        )

    @staticmethod
    def _reasoning_diagnostics(
        message: dict[str, Any],
    ) -> tuple[tuple[str, ...], tuple[str, ...]]:
        present: list[str] = []
        nonempty: list[str] = []
        for name in _REASONING_FIELD_NAMES:
            if name not in message:
                continue
            present.append(name)
            value = message.get(name)
            if isinstance(value, str):
                has_value = bool(value.strip())
            else:
                has_value = value is not None and bool(value)
            if has_value:
                nonempty.append(name)
        return tuple(present), tuple(nonempty)

    @staticmethod
    def _usage_int(usage: dict[str, Any], name: str) -> int:
        try:
            return int(usage.get(name) or 0)
        except (TypeError, ValueError):
            return 0

    async def chat_completion(
        self,
        messages: list[dict],
        max_tokens: int,
        temperature: Optional[float],
    ) -> ChatResult:
        """Send a non-streaming chat completion request."""
        body = self._build_body(messages, max_tokens, temperature, stream=False)
        try:
            response = await self._client.post(self._chat_url, json=body)
        except httpx.HTTPError as e:
            raise self._map_transport_error(e) from e
        if response.status_code != 200:
            raise self._status_error(response.status_code, response.text)
        try:
            data = response.json()
        except ValueError as e:
            raise ExternalEndpointError(
                "External endpoint returned a non-JSON response",
                status="invalid_response",
            ) from e

        try:
            if not isinstance(data, dict):
                raise TypeError("top-level response is not an object")
            choices = data.get("choices")
            if not isinstance(choices, list) or not choices:
                raise KeyError("choices[0]")
            choice = choices[0]
            if not isinstance(choice, dict):
                raise TypeError("choices[0] is not an object")
            message = choice.get("message")
            if not isinstance(message, dict):
                raise KeyError("choices[0].message")
        except (KeyError, TypeError) as e:
            raise ExternalEndpointError(
                f"External endpoint returned an unexpected response shape: {e}",
                status="invalid_response",
            ) from e

        content = message.get("content")
        if content is None:
            text = ""
        elif isinstance(content, str):
            text = content
        else:
            raise ExternalEndpointError(
                "External endpoint returned non-text message.content",
                status="invalid_response",
            )

        finish_reason_value = choice.get("finish_reason")
        finish_reason = (
            str(finish_reason_value) if finish_reason_value is not None else None
        )
        reasoning_present, reasoning_nonempty = self._reasoning_diagnostics(message)
        if finish_reason == "length":
            status = "truncated"
        elif not text.strip():
            status = "empty_content"
        else:
            status = "ok"

        usage = data.get("usage") or {}
        if not isinstance(usage, dict):
            usage = {}
        return ChatResult(
            text=text,
            prompt_tokens=self._usage_int(usage, "prompt_tokens"),
            completion_tokens=self._usage_int(usage, "completion_tokens"),
            finish_reason=finish_reason,
            status=status,
            reasoning_fields_present=reasoning_present,
            reasoning_fields_nonempty=reasoning_nonempty,
        )

    async def stream_chat_completion(
        self,
        messages: list[dict],
        max_tokens: int,
        temperature: Optional[float],
    ) -> StreamStats:
        """Send a streaming chat completion request and collect stats.

        Requires the endpoint to return usage via stream_options
        (include_usage); raises ExternalEndpointError otherwise because
        token counts cannot be measured accurately without it.
        """
        body = self._build_body(messages, max_tokens, temperature, stream=True)
        start_time = time.perf_counter()
        first_content_time: Optional[float] = None
        last_content_time: Optional[float] = None
        usage: Optional[dict] = None
        text_parts: list[str] = []

        try:
            async with self._client.stream(
                "POST", self._chat_url, json=body
            ) as response:
                if response.status_code != 200:
                    error_body = await response.aread()
                    raise self._status_error(
                        response.status_code,
                        error_body.decode("utf-8", errors="replace"),
                    )
                async for line in response.aiter_lines():
                    if not line.startswith("data:"):
                        continue
                    payload = line[len("data:") :].strip()
                    if payload == "[DONE]":
                        break
                    try:
                        chunk = json.loads(payload)
                    except ValueError:
                        continue
                    chunk_usage = chunk.get("usage")
                    if chunk_usage:
                        usage = chunk_usage
                    choices = chunk.get("choices") or []
                    if not choices:
                        continue
                    delta = choices[0].get("delta") or {}
                    content = delta.get("content")
                    if content:
                        text_parts.append(content)
                    has_reasoning = any(
                        delta.get(name) for name in _REASONING_FIELD_NAMES
                    )
                    if content or has_reasoning:
                        now = time.perf_counter()
                        if first_content_time is None:
                            first_content_time = now
                        last_content_time = now
        except httpx.HTTPError as e:
            raise self._map_transport_error(e) from e

        end_time = time.perf_counter()

        if (
            usage is None
            or usage.get("prompt_tokens") is None
            or usage.get("completion_tokens") is None
        ):
            raise ExternalEndpointError(
                "External endpoint does not support stream usage "
                "(stream_options.include_usage); cannot measure token counts"
            )
        content_observed = first_content_time is not None
        if first_content_time is None:
            first_content_time = end_time
        if last_content_time is None:
            last_content_time = end_time
        details = usage.get("prompt_tokens_details") or {}
        return StreamStats(
            prompt_tokens=int(usage["prompt_tokens"]),
            completion_tokens=int(usage["completion_tokens"]),
            cached_tokens=int(details.get("cached_tokens") or 0),
            start_time=start_time,
            first_content_time=first_content_time,
            last_content_time=last_content_time,
            end_time=end_time,
            text="".join(text_parts),
            content_observed=content_observed,
        )


@dataclass
class _AdapterOutput:
    """Minimal GenerationOutput stand-in; eval code only reads .text."""

    text: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
    external_status: str = "ok"
    finish_reason: Optional[str] = None
    reasoning_fields_present: tuple[str, ...] = ()
    reasoning_fields_nonempty: tuple[str, ...] = ()
    error_message: str = ""


class ExternalChatAdapter:
    """Duck-typed engine for accuracy benchmarks against an external API.

    eval/base.py only touches engine.model_type and engine.chat(), so this
    adapter maps both onto ExternalAPIClient. Sampling comes from the
    constructor-injected profile: the temperature/penalty defaults that
    _eval_single injects via setdefault cannot be told apart from
    profile-supplied values, so all sampling kwargs are accepted and
    dropped here — "deterministic" sends temperature 0 and
    "model_settings" sends no sampling params (remote server defaults).
    """

    model_type = None
    is_external_api = True

    def __init__(self, client: ExternalAPIClient, sampling_profile: str):
        self._client = client
        self._sampling_profile = sampling_profile

    async def preflight(self) -> None:
        """Validate final-answer compatibility before a paid evaluation."""
        result = await self._client.chat_completion(
            messages=[{"role": "user", "content": "Reply with exactly: OK"}],
            max_tokens=_PREFLIGHT_MAX_TOKENS,
            temperature=None,
        )
        logger.info(
            "External accuracy preflight: finish_reason=%r "
            "reasoning_fields_present=%s reasoning_fields_nonempty=%s",
            result.finish_reason,
            list(result.reasoning_fields_present),
            list(result.reasoning_fields_nonempty),
        )
        if result.status == "truncated":
            raise ExternalEndpointError(
                "External API preflight was truncated (finish_reason=length)",
                status="truncated",
            )
        if result.status == "empty_content":
            if result.reasoning_fields_nonempty:
                raise ExternalEndpointError(
                    "External API connected, but message.content is empty. "
                    "The model may still be in reasoning mode, or the endpoint "
                    "response format may be incompatible with oMLX.",
                    status="empty_content",
                )
            raise ExternalEndpointError(
                "External API connected, but preflight message.content is empty",
                status="empty_content",
            )
        # Some endpoints inline reasoning into message.content as <think>
        # blocks; drop them before checking, the same way eval scoring does.
        answer = re.sub(
            r"<think>.*?</think>", "", result.text, flags=re.DOTALL
        ).strip()
        # Accept a leading OK with trailing punctuation or extra words
        # (OK. / Okay / OK!) so a compliant endpoint is not rejected over
        # formatting. A genuinely wrong or negated reply (NOT OK) still fails.
        if not answer.upper().startswith("OK"):
            raise ExternalEndpointError(
                "External API preflight response did not start with OK in "
                "message.content",
                status="parse_error",
            )

    async def chat(
        self,
        messages: list[dict],
        max_tokens: int = 256,
        **kwargs: Any,
    ) -> _AdapterOutput:
        temperature = 0.0 if self._sampling_profile == "deterministic" else None
        try:
            result = await self._client.chat_completion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
            )
        except ExternalEndpointError as exc:
            return _AdapterOutput(
                text="",
                external_status=exc.status,
                error_message=str(exc),
            )
        return _AdapterOutput(
            text=result.text,
            prompt_tokens=result.prompt_tokens,
            completion_tokens=result.completion_tokens,
            external_status=result.status,
            finish_reason=result.finish_reason,
            reasoning_fields_present=result.reasoning_fields_present,
            reasoning_fields_nonempty=result.reasoning_fields_nonempty,
        )


# ===== FILE: omlx/admin/hf_downloader.py =====

# SPDX-License-Identifier: Apache-2.0
"""HuggingFace model downloader for oMLX admin panel.

Downloads models from HuggingFace Hub using huggingface_hub's snapshot_download
with directory-size-based progress polling.
"""

import asyncio
import enum
import logging
import shutil
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional
from urllib.parse import urlparse

from huggingface_hub import HfApi, hf_hub_download, snapshot_download
from huggingface_hub.utils import (
    EntryNotFoundError,
    GatedRepoError,
    HfHubHTTPError,
    RepositoryNotFoundError,
)
from huggingface_hub.utils import tqdm as _hf_tqdm

# Private-module import; the pyproject floor (huggingface-hub>=1.19.0)
# guarantees it exists. Re-verify the symbol when bumping the hub version.
from huggingface_hub.utils._xet import abort_xet_session

logger = logging.getLogger(__name__)

# Timeout for HuggingFace API calls (seconds).
# Prevents server from hanging when HF is unreachable.
_HF_API_TIMEOUT = 10

# Seconds with no download progress before considering the download stalled.
_STALL_TIMEOUT = 300

# Cache of (configured_endpoint -> resolved_endpoint) so we only probe each
# endpoint once per process lifetime. Mirrors like hf-mirror.com permanently
# 308-redirect to huggingface.co when accessed from IPs outside their region;
# huggingface_hub does NOT follow those cross-origin 308s during HEAD probes,
# so downloads fail. We resolve the redirect chain upfront and pin HfApi to
# the final origin.
_endpoint_resolution_cache: dict[str, str] = {}


def _resolve_endpoint(endpoint: str) -> str:
    """Follow permanent (301/308) cross-origin redirects on `endpoint`.

    Returns the final origin (scheme://host[:port]) the endpoint resolves to.
    Used to work around `huggingface_hub`'s inability to follow cross-origin
    308 redirects during file-download HEAD probes.

    Probes a known-stable HF API path (`/api/models/gpt2`) with HEAD; if the
    server returns a 301/308 with a Location pointing at a different host,
    the redirected origin is returned (and cached). Network errors fall back
    to the original endpoint.
    """
    endpoint = endpoint.rstrip("/")
    if endpoint in _endpoint_resolution_cache:
        return _endpoint_resolution_cache[endpoint]

    try:
        import httpx
    except ImportError:
        return endpoint

    probe = f"{endpoint}/api/models/gpt2"
    original_host = urlparse(endpoint).netloc
    resolved = endpoint
    try:
        with httpx.Client(follow_redirects=False, timeout=5.0) as client:
            r = client.head(probe)
            # Walk up to 3 permanent hops; stop on first non-permanent status.
            hops = 0
            current_url = probe
            while r.status_code in (301, 308) and "location" in r.headers:
                hops += 1
                if hops > 3:
                    break
                location = r.headers["location"]
                if location.startswith("/"):
                    # Relative redirect — same origin, no rewrite needed.
                    break
                target = urlparse(location)
                if not target.netloc:
                    break
                if target.netloc != original_host:
                    # Cross-origin permanent redirect: rewrite the endpoint.
                    port = f":{target.port}" if target.port else ""
                    resolved = f"{target.scheme}://{target.hostname}{port}"
                    original_host = target.netloc
                current_url = location
                r = client.head(current_url)
    except Exception as e:  # noqa: BLE001 — probe is best-effort
        logger.debug(f"HF endpoint probe failed for {endpoint}: {e}")
        return endpoint

    if resolved != endpoint:
        logger.info(
            f"HuggingFace endpoint {endpoint} permanently redirects to "
            f"{resolved}; using resolved origin for downloads."
        )
    _endpoint_resolution_cache[endpoint] = resolved
    return resolved


class _DownloadCancelled(Exception):
    """Raised inside the download thread to interrupt a cancelled download."""


def _make_cancellable_tqdm(should_cancel: Callable[[], bool]) -> type:
    """Build a tqdm subclass that aborts the download when cancelled.

    huggingface_hub's http_get calls ``progress.update(len(chunk))`` once per
    downloaded chunk (DOWNLOAD_CHUNK_SIZE, 10MB). A running thread can't be
    force-stopped and snapshot_download takes no cancel token, so we cooperate
    from the progress callback: raising here unwinds the download thread
    cleanly within one chunk, releasing its buffers and connection.

    Note: this only interrupts the Python http_get path, which xet-less repos
    and mirror endpoints still use. On the xet path the Rust side defers a
    callback exception until the whole transfer finishes (issue #1322), so
    cancellation there is driven by ``abort_xet_session()`` instead; this
    class is kept as the raise-on-next-chunk backstop for http_get.
    """

    class _CancellableTqdm(_hf_tqdm):
        def update(self, n=1):
            if should_cancel():
                raise _DownloadCancelled()
            return super().update(n)

    return _CancellableTqdm


def _get_hf_api() -> tuple[HfApi, str | None]:
    """Create HfApi instance with configured endpoint.

    Only the admin UI's `huggingface.endpoint` setting is honored here.
    When that's empty, return `HfApi()` with no explicit endpoint so
    `huggingface_hub` falls back to its own resolution (which already
    honors the `HF_ENDPOINT` env var). The configured endpoint, when
    present, is run through `_resolve_endpoint()` to follow permanent
    cross-origin redirects (e.g. hf-mirror.com → huggingface.co from
    non-CN IPs) so downstream HF library code sees a stable origin.

    Returns:
        Tuple of (HfApi instance, endpoint URL or None).
    """
    endpoint: str | None = None
    try:
        from ..settings import get_settings

        endpoint = get_settings().huggingface.endpoint or None
    except (RuntimeError, AttributeError):
        endpoint = None

    if endpoint:
        resolved = _resolve_endpoint(endpoint)
        return HfApi(endpoint=resolved), resolved
    return HfApi(), None


def _list_models_stale_token_fallback(api: HfApi, kwargs: dict) -> tuple[list, bool]:
    """Drain list_models, retrying anonymously when the stored token is rejected.

    huggingface_hub attaches the locally stored credential (HF_TOKEN env var or
    the hf auth login token file) to every request, so a stale token 401s even
    the public model listing (#2276, #2310). Listing needs no auth, so retry
    once with token=False and report the rejected token to the caller.
    """
    try:
        return list(api.list_models(**kwargs)), False
    except HfHubHTTPError as e:
        if e.response is None or e.response.status_code != 401:
            raise
        logger.warning(
            "HF model listing rejected the stored token (401): %s. "
            "Retrying anonymously.",
            e,
        )
        return list(api.list_models(token=False, **kwargs)), True


class DownloadStatus(str, enum.Enum):
    """Status of a download task."""

    PENDING = "pending"
    DOWNLOADING = "downloading"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class DownloadTask:
    """Represents a single model download task."""

    task_id: str
    repo_id: str
    status: DownloadStatus = DownloadStatus.PENDING
    progress: float = 0.0
    total_size: int = 0
    downloaded_size: int = 0
    error: str = ""
    created_at: float = field(default_factory=time.time)
    started_at: float = 0.0
    completed_at: float = 0.0
    retry_count: int = 0

    def to_dict(self) -> dict:
        """Serialize task to a JSON-compatible dict."""
        return {
            "task_id": self.task_id,
            "repo_id": self.repo_id,
            "status": self.status.value,
            "progress": round(self.progress, 1),
            "total_size": self.total_size,
            "downloaded_size": self.downloaded_size,
            "error": self.error,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "retry_count": self.retry_count,
        }


_DTYPE_BYTES = {
    "F64": 8, "F32": 4, "F16": 2, "BF16": 2,
    "I64": 8, "I32": 4, "I16": 2, "I8": 1,
    "U64": 8, "U32": 4, "U16": 2, "U8": 1,
    "BOOL": 1,
}

# Minimum downloads to be included in recommendations.
_MIN_DOWNLOADS = 100


def _calc_safetensors_disk_size(safetensors: dict) -> int:
    """Calculate actual disk size in bytes from safetensors parameters.

    safetensors.total is the parameter count, not bytes.
    We need to multiply each dtype's parameter count by its byte width.
    """
    params = safetensors.get("parameters", {})
    if not params:
        return 0
    return sum(count * _DTYPE_BYTES.get(dtype, 1) for dtype, count in params.items())


def _format_model_size(size_bytes: int) -> str:
    """Format model size in bytes to a human-readable string."""
    if size_bytes < 1024**2:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024**3:
        return f"{size_bytes / 1024**2:.1f} MB"
    else:
        return f"{size_bytes / 1024**3:.1f} GB"


def _format_param_count(total_params: int) -> str:
    """Format parameter count to a human-readable string (e.g., 7.0B, 13.0B)."""
    if total_params >= 1e12:
        return f"{total_params / 1e12:.1f}T"
    if total_params >= 1e9:
        return f"{total_params / 1e9:.1f}B"
    if total_params >= 1e6:
        return f"{total_params / 1e6:.1f}M"
    return str(total_params)


def _get_param_count(safetensors: dict) -> int:
    """Get total parameter count from safetensors metadata."""
    params = safetensors.get("parameters", {})
    if not params:
        return 0
    return sum(params.values())


# HF API sort field mapping for search.
_SORT_MAP = {
    "trending": "trendingScore",
    "downloads": "downloads",
    "created": "createdAt",
    "updated": "lastModified",
    "most_params": "downloads",  # fetch by downloads, re-sort in Python
    "least_params": "downloads",  # fetch by downloads, re-sort in Python
    "largest": "downloads",  # fetch by downloads, re-sort by size in Python
    "smallest": "downloads",  # fetch by downloads, re-sort by size in Python
}


class HFDownloader:
    """Manages HuggingFace model downloads with progress tracking.

    Uses huggingface_hub.snapshot_download() for actual downloads and polls
    the target directory size to estimate progress.

    Args:
        model_dir: Directory where downloaded models are stored.
        on_complete: Async callback invoked when a download completes successfully.
    """

    @staticmethod
    async def get_recommended_models(
        max_memory_bytes: int,
        limit: int = 60,
        result_limit: int = 50,
        mlx_only: bool = True,
    ) -> dict:
        """Fetch trending and popular models that fit in memory.

        Queries HuggingFace Hub for models, optionally restricted to
        mlx-community. Filtered by system memory capacity.

        Args:
            max_memory_bytes: Maximum model size in bytes (typically system memory).
            limit: Number of models to fetch per category from HF API.
            result_limit: Maximum number of models to return per category.
            mlx_only: If True, restrict to mlx-community author.

        Returns:
            Dict with 'trending' and 'popular' lists, plus 'hf_token_invalid'
            set when the stored HF token was rejected and the listing was
            fetched anonymously instead.
        """
        api, _endpoint = _get_hf_api()

        async def _fetch(sort: str) -> tuple[list[dict], bool]:
            kwargs = {
                "sort": sort,
                "limit": limit,
                "expand": ["safetensors", "downloads", "likes", "trendingScore"],
            }
            if mlx_only:
                kwargs["author"] = "mlx-community"
            # list_models returns a lazy generator; drain it inside the worker
            # thread so the paginated HTTP calls never block the event loop.
            models, token_rejected = await asyncio.wait_for(
                asyncio.to_thread(_list_models_stale_token_fallback, api, kwargs),
                timeout=_HF_API_TIMEOUT,
            )
            results = []
            for m in models:
                if not m.safetensors or not m.safetensors.get("parameters"):
                    continue
                downloads = m.downloads or 0
                if downloads < _MIN_DOWNLOADS:
                    continue
                size = _calc_safetensors_disk_size(m.safetensors)
                if size <= 0 or size > max_memory_bytes:
                    continue
                params = _get_param_count(m.safetensors)
                results.append(
                    {
                        "repo_id": m.id,
                        "name": m.id.split("/")[-1],
                        "downloads": downloads,
                        "likes": m.likes or 0,
                        "trending_score": m.trending_score or 0,
                        "size": size,
                        "size_formatted": _format_model_size(size),
                        "params": params if params > 0 else None,
                        "params_formatted": (
                            _format_param_count(params) if params > 0 else None
                        ),
                    }
                )
            return results, token_rejected

        (trending, trending_rejected), (popular, popular_rejected) = (
            await asyncio.gather(
                _fetch("trendingScore"),
                _fetch("downloads"),
            )
        )

        return {
            "trending": trending[:result_limit],
            "popular": popular[:result_limit],
            "hf_token_invalid": trending_rejected or popular_rejected,
        }

    @staticmethod
    async def search_models(
        query: str,
        sort: str = "trending",
        limit: int = 100,
        mlx_only: bool = True,
        # Filtering options
        min_params: Optional[int] = None,
        max_params: Optional[int] = None,
        min_size: Optional[int] = None,
        max_size: Optional[int] = None,
        # Sorting options
        sort_by_size: bool = False,
        sort_ascending: bool = False,
    ) -> dict:
        """Search HuggingFace models by query string with filtering and sorting.

        When mlx_only is True, results are restricted to the MLX library
        (same as https://huggingface.co/models?library=mlx).

        Args:
            query: Search query string.
            sort: Sort order (trending/downloads/created/updated/most_params/least_params/largest/smallest).
            limit: Maximum number of results to return.
            mlx_only: If True, restrict to MLX library models only.
            min_params: Minimum parameter count filter.
            max_params: Maximum parameter count filter.
            min_size: Minimum model size in bytes filter.
            max_size: Maximum model size in bytes filter.
            sort_by_size: Sort results by size instead of default sort.
            sort_ascending: Sort in ascending order (for size/params sorting).

        Returns:
            Dict with 'models' list and 'total' count, plus 'hf_token_invalid'
            set when the stored HF token was rejected and the listing was
            fetched anonymously instead.
        """
        api, _endpoint = _get_hf_api()

        # Determine base sort - for Python-side sorting, we fetch by downloads
        # which tends to return more results, then sort in Python
        if sort in ("most_params", "least_params", "largest", "smallest"):
            base_sort = "downloads"
        else:
            base_sort = _SORT_MAP.get(sort, "trendingScore")

        kwargs = {
            "search": query,
            "sort": base_sort,
            "limit": limit,
            "expand": ["safetensors", "downloads", "likes", "trendingScore"],
        }
        if mlx_only:
            kwargs["filter"] = "mlx"

        # list_models returns a lazy generator; drain it inside the worker
        # thread so the paginated HTTP calls never block the event loop.
        models, token_rejected = await asyncio.wait_for(
            asyncio.to_thread(_list_models_stale_token_fallback, api, kwargs),
            timeout=_HF_API_TIMEOUT,
        )

        results = []
        for m in models:
            params = None
            params_formatted = None
            size = 0

            if m.safetensors and m.safetensors.get("parameters"):
                params = _get_param_count(m.safetensors)
                params_formatted = _format_param_count(params) if params > 0 else None
                size = _calc_safetensors_disk_size(m.safetensors)
                if params and params <= 0:
                    params = None

            # Apply filters
            if min_params is not None and (params is None or params < min_params):
                continue
            if max_params is not None and (params is None or params > max_params):
                continue
            if min_size is not None and size < min_size:
                continue
            if max_size is not None and size > max_size:
                continue

            results.append(
                {
                    "repo_id": m.id,
                    "name": m.id,
                    "downloads": m.downloads or 0,
                    "likes": m.likes or 0,
                    "trending_score": m.trending_score or 0,
                    "size": size,
                    "size_formatted": _format_model_size(size) if size > 0 else "",
                    "params": params,
                    "params_formatted": params_formatted,
                }
            )

        # Apply Python-side sorting
        if sort == "most_params":
            results.sort(key=lambda x: x["params"] or 0, reverse=True)
        elif sort == "least_params":
            results.sort(key=lambda x: x["params"] or 0)
        elif sort in ("largest", "smallest") or sort_by_size:
            # Sort by size, putting unknown-size entries at the end
            results.sort(
                key=lambda x: x["size"] if x["size"] > 0 else -1,
                reverse=(sort == "largest" or (sort_by_size and not sort_ascending)),
            )
        # Otherwise, keep original HF API ordering (trending, downloads, created, updated)

        return {
            "models": results[:limit],
            "total": len(results),
            "hf_token_invalid": token_rejected,
        }

    @staticmethod
    async def get_model_info(repo_id: str) -> dict:
        """Fetch detailed model information from HuggingFace.

        Args:
            repo_id: HuggingFace repository ID (e.g., "mlx-community/Llama-3-8B-4bit").

        Returns:
            Dict with model details including description, files, tags, etc.
        """
        api, endpoint = _get_hf_api()
        info = await asyncio.wait_for(
            asyncio.to_thread(
                api.model_info,
                repo_id,
                files_metadata=True,
            ),
            timeout=_HF_API_TIMEOUT,
        )

        # Extract file list with sizes
        files = []
        if info.siblings:
            for s in info.siblings:
                files.append(
                    {
                        "name": s.rfilename,
                        "size": s.size or 0,
                        "size_formatted": (
                            _format_model_size(s.size) if s.size else ""
                        ),
                    }
                )

        # Detect LoRA/adapter repos (adapter_config.json is peft standard)
        is_adapter = any(f["name"] == "adapter_config.json" for f in files)

        # Extract params and size from safetensors
        params = None
        params_formatted = None
        size = 0
        safetensors = getattr(info, "safetensors", None)
        if safetensors:
            st_dict = dict(safetensors) if not isinstance(safetensors, dict) else safetensors
            if st_dict.get("parameters"):
                params = _get_param_count(st_dict)
                params_formatted = _format_param_count(params) if params > 0 else None
                size = _calc_safetensors_disk_size(st_dict)

        # Fetch model card (README.md) content
        model_card = ""
        try:
            card_path = await asyncio.wait_for(
                asyncio.to_thread(
                    hf_hub_download,
                    repo_id=repo_id,
                    filename="README.md",
                    endpoint=endpoint,
                ),
                timeout=_HF_API_TIMEOUT,
            )
            if card_path:
                card_text = Path(card_path).read_text(encoding="utf-8")
                # Strip YAML front matter (between --- markers)
                if card_text.startswith("---"):
                    end = card_text.find("---", 3)
                    if end != -1:
                        card_text = card_text[end + 3:].strip()
                model_card = card_text
        except Exception:
            pass  # README not available

        return {
            "repo_id": info.id,
            "name": info.id,
            "model_card": model_card,
            "description": "",  # kept for backward compat
            "files": files,
            "tags": list(info.tags) if info.tags else [],
            "pipeline_tag": info.pipeline_tag or "",
            "params": params,
            "params_formatted": params_formatted,
            "size": size,
            "size_formatted": _format_model_size(size) if size > 0 else "",
            "downloads": info.downloads or 0,
            "likes": info.likes or 0,
            "created_at": info.created_at.isoformat() if info.created_at else "",
            "updated_at": (
                info.last_modified.isoformat() if info.last_modified else ""
            ),
            "is_adapter": is_adapter,
        }

    def __init__(
        self,
        model_dir: str,
        on_complete: Optional[Callable] = None,
    ):
        self._model_dir = Path(model_dir)
        self._tasks: dict[str, DownloadTask] = {}
        self._active_tasks: dict[str, asyncio.Task] = {}
        self._progress_tasks: dict[str, asyncio.Task] = {}
        self._on_complete = on_complete
        self._cancelled: set[str] = set()
        self._download_sem = asyncio.Semaphore(1)

    @property
    def model_dir(self) -> Path:
        return self._model_dir

    def update_model_dir(self, new_dir: str) -> None:
        """Update the model directory path."""
        self._model_dir = Path(new_dir)

    async def start_download(
        self, repo_id: str, hf_token: str = ""
    ) -> DownloadTask:
        """Start downloading a model from HuggingFace.

        Args:
            repo_id: HuggingFace repository ID (e.g., "mlx-community/Llama-3-8B-4bit").
            hf_token: Optional HuggingFace token for gated models.

        Returns:
            The created DownloadTask.

        Raises:
            ValueError: If repo_id format is invalid or download is already queued.
        """
        repo_id = repo_id.strip()
        if "/" not in repo_id or len(repo_id.split("/")) != 2:
            raise ValueError(
                f"Invalid repository ID: '{repo_id}'. "
                "Expected format: 'owner/model' (e.g., 'mlx-community/Llama-3-8B-4bit')"
            )

        # Check for duplicate active downloads
        for task in self._tasks.values():
            if task.repo_id == repo_id and task.status in (
                DownloadStatus.PENDING,
                DownloadStatus.DOWNLOADING,
            ):
                raise ValueError(
                    f"Download for '{repo_id}' is already in progress"
                )

        task_id = str(uuid.uuid4())
        task = DownloadTask(task_id=task_id, repo_id=repo_id)
        self._tasks[task_id] = task

        # Start download in background
        self._active_tasks[task_id] = asyncio.create_task(
            self._run_download(task_id, hf_token)
        )

        logger.info(f"Download queued: {repo_id} (task_id={task_id})")
        return task

    async def cancel_download(self, task_id: str) -> bool:
        """Cancel an active download.

        Args:
            task_id: The task ID to cancel.

        Returns:
            True if the task was found and cancelled.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status not in (DownloadStatus.PENDING, DownloadStatus.DOWNLOADING):
            return False

        was_downloading = task.status == DownloadStatus.DOWNLOADING

        # Mark as cancelled
        self._cancelled.add(task_id)
        task.status = DownloadStatus.CANCELLED

        # A task in DOWNLOADING owns the download semaphore, so the in-flight
        # xet transfer is necessarily this one; aborting the (global) session
        # makes its snapshot_download thread unwind immediately. Pending tasks
        # must not abort, that would kill another task's transfer. The next
        # download lazily creates a fresh session.
        if was_downloading:
            abort_xet_session()

        # Stop progress polling
        progress_task = self._progress_tasks.pop(task_id, None)
        if progress_task and not progress_task.done():
            progress_task.cancel()

        # Cancel the download task
        active_task = self._active_tasks.pop(task_id, None)
        if active_task and not active_task.done():
            active_task.cancel()

        logger.info(f"Download cancelled: {task.repo_id} (task_id={task_id})")
        return True

    def remove_task(self, task_id: str) -> bool:
        """Remove a completed, failed, or cancelled task from the list.

        Args:
            task_id: The task ID to remove.

        Returns:
            True if the task was found and removed.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status in (DownloadStatus.PENDING, DownloadStatus.DOWNLOADING):
            return False

        del self._tasks[task_id]
        self._cancelled.discard(task_id)
        return True

    async def retry_download(
        self, task_id: str, hf_token: str = ""
    ) -> DownloadTask:
        """Retry a failed or cancelled download, resuming from existing files.

        Finalized shards are preserved on disk so snapshot_download will
        automatically skip already-completed files.

        Args:
            task_id: The task ID of the failed/cancelled download.
            hf_token: Optional HuggingFace token for gated models.

        Returns:
            The new DownloadTask.

        Raises:
            ValueError: If task not found or not in retryable state.
        """
        old_task = self._tasks.get(task_id)
        if old_task is None:
            raise ValueError(f"Task not found: {task_id}")

        if old_task.status not in (DownloadStatus.FAILED, DownloadStatus.CANCELLED):
            raise ValueError(
                f"Task {task_id} is not retryable (status: {old_task.status.value})"
            )

        repo_id = old_task.repo_id
        old_retry_count = old_task.retry_count

        # Remove old task entry
        del self._tasks[task_id]
        self._cancelled.discard(task_id)

        # Start fresh download (snapshot_download resumes from existing files)
        new_task = await self.start_download(repo_id, hf_token)
        new_task.retry_count = old_retry_count + 1
        return new_task

    def get_tasks(self) -> list[dict]:
        """Return all tasks as serializable dicts, ordered by creation time."""
        return [
            task.to_dict()
            for task in sorted(self._tasks.values(), key=lambda t: t.created_at)
        ]

    async def shutdown(self) -> None:
        """Cancel all active downloads and clean up."""
        # Cancel all progress polling tasks
        for task_id, progress_task in list(self._progress_tasks.items()):
            if not progress_task.done():
                progress_task.cancel()
        self._progress_tasks.clear()

        # Cancel all active download tasks. Mark cancelled first so an
        # in-flight snapshot_download thread aborts via its progress callback;
        # active_task.cancel() only unblocks tasks still waiting on the semaphore.
        for task_id, active_task in list(self._active_tasks.items()):
            self._cancelled.add(task_id)
            if not active_task.done():
                active_task.cancel()
                task = self._tasks.get(task_id)
                if task and task.status == DownloadStatus.DOWNLOADING:
                    task.status = DownloadStatus.CANCELLED
        self._active_tasks.clear()

        # Reap any in-flight xet transfer thread (no-op without a session).
        abort_xet_session()

        logger.info("HF Downloader shut down")

    async def _run_download(self, task_id: str, hf_token: str) -> None:
        """Execute a download task.

        Waits for the download semaphore (only one download runs at a time),
        then fetches repo info for total size and runs snapshot_download in a
        thread while polling the target directory for progress updates.
        """
        task = self._tasks[task_id]

        try:
            async with self._download_sem:
                # Check if cancelled while waiting in queue
                if task_id in self._cancelled:
                    return

                task.status = DownloadStatus.DOWNLOADING
                task.started_at = time.time()

                # Preserve {owner}/{model} layout to match other tools
                # (LMStudio, huggingface-cli) and avoid duplicate downloads
                # when sharing a model directory.
                target_dir = self._model_dir / task.repo_id

                api, endpoint = _get_hf_api()

                # Skip pytorch format when safetensors exist to
                # avoid downloading redundant weight files.
                ignore_patterns = None
                st_estimate = 0
                try:
                    model_info = await asyncio.wait_for(
                        asyncio.to_thread(
                            api.model_info,
                            task.repo_id,
                            token=hf_token or None,
                            expand=["safetensors"],
                        ),
                        timeout=_HF_API_TIMEOUT,
                    )
                    if model_info.safetensors and model_info.safetensors.get(
                        "parameters"
                    ):
                        ignore_patterns = [
                            "*.bin",
                            "original/**",
                            "consolidated.*.pth",
                        ]
                        # Computed inside this try so malformed metadata
                        # (non-int counts) degrades to no estimate instead
                        # of failing the download from the dry-run handler.
                        st_estimate = _calc_safetensors_disk_size(
                            model_info.safetensors
                        )
                except Exception as e:
                    logger.warning(
                        f"Could not fetch repo info for {task.repo_id}: {e}"
                    )

                dl_kwargs: dict = {
                    "repo_id": task.repo_id,
                    "local_dir": str(target_dir),
                    "token": hf_token or None,
                    "endpoint": endpoint,
                    "etag_timeout": 30,
                }
                if ignore_patterns:
                    dl_kwargs["ignore_patterns"] = ignore_patterns

                # Get accurate total size via dry run so the progress
                # denominator matches what will actually be downloaded.
                size_estimated = False
                try:
                    dry_result = await asyncio.wait_for(
                        asyncio.to_thread(
                            snapshot_download,
                            **dl_kwargs,
                            dry_run=True,
                        ),
                        timeout=30,
                    )
                    task.total_size = sum(f.file_size for f in dry_result)
                except Exception as e:
                    if st_estimate:
                        task.total_size = st_estimate
                        size_estimated = True
                        detail = "Estimated total size from safetensors metadata."
                    else:
                        detail = "Progress estimation will be unavailable."
                    logger.warning(
                        f"Dry run failed for {task.repo_id}: {e}. {detail}"
                    )

                # Start progress polling
                self._progress_tasks[task_id] = asyncio.create_task(
                    self._poll_progress(task_id, target_dir)
                )

                # Run snapshot_download in a thread (blocking call). Cancel
                # reaches the thread two ways: the cancellable tqdm raises on
                # the next chunk of the http_get path, and abort_xet_session()
                # (called by cancel/stall/shutdown) unwinds the xet path with
                # a RuntimeError (the thread itself can't be force-killed).
                await asyncio.to_thread(
                    snapshot_download,
                    **dl_kwargs,
                    tqdm_class=_make_cancellable_tqdm(
                        lambda: task_id in self._cancelled
                    ),
                )

                # Check if cancelled while downloading
                if task_id in self._cancelled:
                    return

                # Success
                task.status = DownloadStatus.COMPLETED
                task.progress = 100.0
                if size_estimated or not task.total_size:
                    # The estimate was only a progress denominator; report
                    # the measured on-disk size once the download is done.
                    task.downloaded_size = self._get_dir_size(target_dir)
                else:
                    task.downloaded_size = task.total_size
                task.completed_at = time.time()

                logger.info(
                    f"Download completed: {task.repo_id} -> {target_dir} "
                    f"({time.time() - task.started_at:.1f}s)"
                )

                # Trigger model pool refresh
                if self._on_complete:
                    try:
                        await self._on_complete()
                    except Exception as e:
                        logger.error(
                            f"Error in download completion callback: {e}"
                        )

        except (_DownloadCancelled, asyncio.CancelledError):
            if task.status not in (
                DownloadStatus.CANCELLED,
                DownloadStatus.FAILED,
            ):
                task.status = DownloadStatus.CANCELLED
            try:
                self._cleanup_partial(task)
            except Exception as e:
                logger.error(
                    f"Failed to clean up cancelled download {task.repo_id}: {e}"
                )
        except RepositoryNotFoundError:
            task.status = DownloadStatus.FAILED
            task.error = (
                f"Repository not found: {task.repo_id}. "
                "This may be a gated model that requires HuggingFace authentication."
            )
            logger.error(f"Repository not found: {task.repo_id}")
        except GatedRepoError:
            task.status = DownloadStatus.FAILED
            task.error = (
                f"Repository '{task.repo_id}' is gated. "
                "Please provide a valid HF token with access."
            )
            logger.error(f"Gated repo access denied: {task.repo_id}")
        except Exception as e:
            # Skip when already cancelled (the xet abort surfaces here as a
            # RuntimeError) or already FAILED by the stall detector, whose
            # error message would otherwise be clobbered by the abort error.
            if (
                task_id not in self._cancelled
                and task.status != DownloadStatus.FAILED
            ):
                task.status = DownloadStatus.FAILED
                task.error = str(e)
                logger.error(f"Download failed for {task.repo_id}: {e}")
        finally:
            # Stop progress polling
            progress_task = self._progress_tasks.pop(task_id, None)
            if progress_task and not progress_task.done():
                progress_task.cancel()

            # Remove from active tasks
            self._active_tasks.pop(task_id, None)

    async def _poll_progress(self, task_id: str, target_dir: Path) -> None:
        """Poll the target directory to estimate download progress.

        Uses both directory size and file modification times to detect
        activity. huggingface_hub pre-allocates large files and fills them
        in, so size alone may not change for extended periods. File mtimes
        are updated on each write syscall and serve as a more reliable
        liveness signal.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return

        last_size = 0
        last_activity_at = time.time()

        try:
            while task.status == DownloadStatus.DOWNLOADING:
                await asyncio.sleep(2)

                if task.status != DownloadStatus.DOWNLOADING:
                    break

                current_size = self._get_dir_size(target_dir)
                task.downloaded_size = current_size

                if task.total_size > 0:
                    # Cap at 99% until snapshot_download confirms completion
                    task.progress = min(
                        (current_size / task.total_size) * 100, 99.0
                    )

                # Activity detection: size change OR file mtime change
                if current_size != last_size:
                    last_size = current_size
                    last_activity_at = time.time()
                else:
                    latest_mtime = self._get_latest_mtime(target_dir)
                    if latest_mtime > last_activity_at:
                        last_activity_at = latest_mtime

                # Stall detection
                if (
                    current_size > 0
                    and (time.time() - last_activity_at) > _STALL_TIMEOUT
                ):
                    task.status = DownloadStatus.FAILED
                    task.error = (
                        f"Download stalled: no progress for {_STALL_TIMEOUT}s. "
                        "Try retrying the download."
                    )
                    logger.warning(
                        f"Download stalled for {task.repo_id} "
                        f"(task_id={task_id})"
                    )
                    # Cancel the snapshot_download thread. The task cancel
                    # only unblocks the awaiting coroutine; aborting the xet
                    # session is what actually reaps a wedged transfer thread.
                    active_task = self._active_tasks.get(task_id)
                    if active_task and not active_task.done():
                        active_task.cancel()
                    abort_xet_session()
                    break
        except asyncio.CancelledError:
            pass

    @staticmethod
    def _get_latest_mtime(path: Path) -> float:
        """Return the most recent modification time of any file in a directory."""
        if not path.exists():
            return 0.0
        latest = 0.0
        try:
            for f in path.rglob("*"):
                if f.is_file():
                    try:
                        mt = f.stat().st_mtime
                        if mt > latest:
                            latest = mt
                    except OSError:
                        pass
        except OSError:
            pass
        return latest

    @staticmethod
    def _get_dir_size(path: Path) -> int:
        """Calculate total size of all files in a directory."""
        if not path.exists():
            return 0
        total = 0
        try:
            for f in path.rglob("*"):
                if f.is_file():
                    try:
                        total += f.stat().st_size
                    except OSError:
                        pass
        except OSError:
            pass
        return total

    def _cleanup_partial(self, task: DownloadTask) -> None:
        """Remove in-progress shards while keeping finalized files for resume.

        Hub stages partial downloads inside a hidden ``._____temp`` directory
        and only renames a shard into the target on completion. Wiping the
        whole target dir would also nuke shards the user has already paid
        for; finalized files are visible in the file browser, so users can
        keep them for auto-resume on retry or remove them themselves.
        """
        target_dir = self._model_dir / task.repo_id
        temp_dir = target_dir / "._____temp"
        if temp_dir.exists():
            try:
                shutil.rmtree(temp_dir)
                logger.info(f"Cleaned up in-progress shards: {temp_dir}")
            except Exception as e:
                logger.error(f"Failed to clean up {temp_dir}: {e}")


# ===== FILE: omlx/admin/hf_uploader.py =====

# SPDX-License-Identifier: Apache-2.0
"""HuggingFace model uploader for oMLX admin panel.

Uploads oQ-quantized models to HuggingFace Hub with queue-based sequential
processing, following the same pattern as hf_downloader.py.
"""

import asyncio
import enum
import json
import logging
import shutil
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


def _format_size(size_bytes: int) -> str:
    """Format size in bytes to human-readable string."""
    if size_bytes < 1024**2:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024**3:
        return f"{size_bytes / 1024**2:.1f} MB"
    else:
        return f"{size_bytes / 1024**3:.1f} GB"


def _has_meaningful_readme(path: Path) -> bool:
    """Check if a README.md exists and has content beyond YAML frontmatter.

    Returns False if the file doesn't exist, is empty, or contains only
    YAML frontmatter (e.g. mlx-lm's default stub).
    """
    readme = path / "README.md"
    if not readme.exists():
        return False
    try:
        text = readme.read_text(encoding="utf-8").strip()
    except Exception:
        return False
    if not text:
        return False
    # Strip YAML frontmatter and check if anything remains
    if text.startswith("---"):
        parts = text.split("---", 2)
        # parts[0] is empty (before first ---), parts[1] is frontmatter
        if len(parts) >= 3:
            body = parts[2].strip()
            return len(body) > 0
        # Only opening --- or unclosed frontmatter
        return False
    return True


def _is_oq_model(name: str) -> bool:
    """Check if a model name indicates an oQ-quantized model.

    Any folder name containing 'oQ' (case-sensitive) is treated as an oQ model,
    e.g. 'Qwen3.5-122B-oQ4', 'Llama-3B-oQ4e', 'Qwen3.6-27B-oQ3.5e'.
    """
    return "oQ" in name


def _generate_model_card(
    model_name: str, config: dict, redownload_notice: bool = False,
) -> str:
    """Generate a minimal HuggingFace model card for an oQ model."""
    from omlx._version import __version__

    model_type = config.get("model_type", "unknown")
    quant = config.get("quantization", {})
    bits = quant.get("bits", "?")
    group_size = quant.get("group_size", "?")

    notice = ""
    if redownload_notice:
        from datetime import date

        today = date.today().strftime("%Y-%m-%d")
        notice = f"""> [!IMPORTANT]
> This quantization was uploaded on **{today}** and replaces a previous version.
> If you downloaded this model before this date, please re-download for the updated weights.

"""

    return f"""---
library_name: mlx
tags:
- mlx
- oq
- quantized
---

{notice}# {model_name}

This model was quantized using [oQ](https://github.com/jundot/omlx) (oMLX v{__version__}) mixed-precision quantization.

## Quantization details

- **Model type**: {model_type}
- **Bits**: {bits}
- **Group size**: {group_size}
- **Format**: MLX safetensors
"""


class UploadStatus(str, enum.Enum):
    """Status of an upload task."""

    PENDING = "pending"
    UPLOADING = "uploading"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


_ACTIVE_STATUSES = {UploadStatus.PENDING, UploadStatus.UPLOADING}


@dataclass
class UploadTask:
    """Represents a single model upload task."""

    task_id: str
    model_name: str
    model_path: str
    repo_id: str
    status: UploadStatus = UploadStatus.PENDING
    progress: float = 0.0
    error: str = ""
    created_at: float = field(default_factory=time.time)
    started_at: float = 0.0
    completed_at: float = 0.0
    total_size: int = 0
    repo_url: str = ""

    def to_dict(self) -> dict:
        """Serialize task to a JSON-compatible dict."""
        return {
            "task_id": self.task_id,
            "model_name": self.model_name,
            "model_path": self.model_path,
            "repo_id": self.repo_id,
            "status": self.status.value,
            "progress": round(self.progress, 1),
            "error": self.error,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "total_size": self.total_size,
            "total_size_formatted": _format_size(self.total_size) if self.total_size else "",
            "repo_url": self.repo_url,
        }


class HFUploader:
    """Manages HuggingFace model uploads with queue-based sequential processing.

    Uses huggingface_hub's upload_folder() with a semaphore to ensure only
    one upload runs at a time. Multiple uploads can be queued.

    Args:
        model_dirs: List of model directory paths to scan for oQ models.
    """

    def __init__(self, model_dirs: list[str]):
        self._model_dirs = [Path(d) for d in model_dirs]
        self._tasks: dict[str, UploadTask] = {}
        self._active_tasks: dict[str, asyncio.Task] = {}
        self._cancelled: set[str] = set()
        self._upload_sem = asyncio.Semaphore(1)

    def update_model_dirs(self, model_dirs: list[str]) -> None:
        """Update model directory paths."""
        self._model_dirs = [Path(d) for d in model_dirs]

    @staticmethod
    async def validate_token(token: str) -> dict:
        """Validate a HuggingFace token and return user info.

        Args:
            token: HuggingFace write-access token.

        Returns:
            Dict with 'username' and 'orgs' list.

        Raises:
            ValueError: If the token is invalid or lacks write access.
        """
        from huggingface_hub import HfApi

        try:
            api = HfApi()
            info = await asyncio.to_thread(api.whoami, token=token)
        except Exception as e:
            raise ValueError(f"Invalid token: {e}")

        username = info.get("name", "")
        orgs = [
            {"name": org.get("name", "")}
            for org in info.get("orgs", [])
            if org.get("name")
        ]

        # Check for write access
        auth = info.get("auth", {})
        access_token = auth.get("accessToken", {})
        role = access_token.get("role", "")
        if role == "read":
            raise ValueError(
                "Token has read-only access. A write token is required for uploads."
            )

        return {"username": username, "orgs": orgs}

    async def list_oq_models(self) -> list[dict]:
        """Scan model directories and return oQ-quantized models.

        Returns:
            List of dicts with model name, path, size info.
        """

        def _scan() -> list[dict]:
            models = []
            seen: set[str] = set()

            for model_dir in self._model_dirs:
                if not model_dir.exists():
                    continue
                for subdir in sorted(model_dir.iterdir()):
                    if not subdir.is_dir():
                        continue
                    candidates = []
                    if (subdir / "config.json").exists():
                        candidates.append(subdir)
                    else:
                        for child in sorted(subdir.iterdir()):
                            if child.is_dir() and (child / "config.json").exists():
                                candidates.append(child)

                    for path in candidates:
                        if path.name in seen:
                            continue
                        seen.add(path.name)
                        if not _is_oq_model(path.name):
                            continue
                        try:
                            size = sum(
                                f.stat().st_size
                                for f in path.glob("*.safetensors")
                            )
                            if size == 0:
                                continue
                            models.append({
                                "name": path.name,
                                "path": str(path),
                                "size": size,
                                "size_formatted": _format_size(size),
                            })
                        except Exception:
                            continue
            return models

        return await asyncio.to_thread(_scan)

    async def list_all_models(self) -> list[dict]:
        """Scan model directories and return all models (for README source selection).

        Returns:
            List of dicts with model name and path.
        """

        def _scan() -> list[dict]:
            models = []
            seen: set[str] = set()

            for model_dir in self._model_dirs:
                if not model_dir.exists():
                    continue
                for subdir in sorted(model_dir.iterdir()):
                    if not subdir.is_dir():
                        continue
                    candidates = []
                    if (subdir / "config.json").exists():
                        candidates.append(subdir)
                    else:
                        for child in sorted(subdir.iterdir()):
                            if child.is_dir() and (child / "config.json").exists():
                                candidates.append(child)

                    for path in candidates:
                        if path.name in seen:
                            continue
                        seen.add(path.name)
                        has_readme = _has_meaningful_readme(path)
                        models.append({
                            "name": path.name,
                            "path": str(path),
                            "has_readme": has_readme,
                        })
            return models

        return await asyncio.to_thread(_scan)

    async def start_upload(
        self,
        model_path: str,
        repo_id: str,
        token: str,
        readme_source_path: str = "",
        auto_readme: bool = True,
        redownload_notice: bool = False,
        private: bool = False,
    ) -> UploadTask:
        """Queue a model upload to HuggingFace Hub.

        Args:
            model_path: Local path to the oQ model directory.
            repo_id: Target HuggingFace repository ID (e.g., 'user/model-oQ4').
            token: HuggingFace write token.
            readme_source_path: Optional path to model whose README.md to copy.
            auto_readme: If True and no readme_source_path, generate a basic README.
            private: If True, create a private repository.

        Returns:
            The created UploadTask.

        Raises:
            ValueError: If model path is invalid or upload is already queued.
        """
        source = Path(model_path)
        if not source.exists() or not source.is_dir():
            raise ValueError(f"Model directory not found: {model_path}")

        if not (source / "config.json").exists():
            raise ValueError(f"Not a valid model directory (no config.json): {model_path}")

        repo_id = repo_id.strip()
        if "/" not in repo_id or len(repo_id.split("/")) != 2:
            raise ValueError(
                f"Invalid repository ID: '{repo_id}'. "
                "Expected format: 'owner/model' (e.g., 'user/Llama-3B-oQ4')"
            )

        # Check for duplicate active uploads
        for task in self._tasks.values():
            if task.repo_id == repo_id and task.status in _ACTIVE_STATUSES:
                raise ValueError(
                    f"Upload to '{repo_id}' is already in progress"
                )

        model_name = source.name
        total_size = sum(
            f.stat().st_size for f in source.rglob("*") if f.is_file()
        )

        task_id = str(uuid.uuid4())
        task = UploadTask(
            task_id=task_id,
            model_name=model_name,
            model_path=model_path,
            repo_id=repo_id,
            total_size=total_size,
        )
        self._tasks[task_id] = task

        self._active_tasks[task_id] = asyncio.create_task(
            self._run_upload(task_id, token, readme_source_path, auto_readme, redownload_notice, private)
        )

        logger.info(f"Upload queued: {model_name} -> {repo_id} (task_id={task_id})")
        return task

    async def cancel_upload(self, task_id: str) -> bool:
        """Cancel an active or pending upload.

        Args:
            task_id: The task ID to cancel.

        Returns:
            True if the task was found and cancelled.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status not in _ACTIVE_STATUSES:
            return False

        self._cancelled.add(task_id)
        task.status = UploadStatus.CANCELLED

        active_task = self._active_tasks.pop(task_id, None)
        if active_task and not active_task.done():
            active_task.cancel()

        logger.info(f"Upload cancelled: {task.model_name} (task_id={task_id})")
        return True

    def remove_task(self, task_id: str) -> bool:
        """Remove a completed, failed, or cancelled task from the list.

        Args:
            task_id: The task ID to remove.

        Returns:
            True if the task was found and removed.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status in _ACTIVE_STATUSES:
            return False

        del self._tasks[task_id]
        self._cancelled.discard(task_id)
        return True

    def get_tasks(self) -> list[dict]:
        """Return all tasks as serializable dicts, ordered by creation time."""
        return [
            task.to_dict()
            for task in sorted(self._tasks.values(), key=lambda t: t.created_at)
        ]

    async def shutdown(self) -> None:
        """Cancel all active uploads and clean up."""
        for task_id, active_task in list(self._active_tasks.items()):
            if not active_task.done():
                active_task.cancel()
                task = self._tasks.get(task_id)
                if task and task.status == UploadStatus.UPLOADING:
                    task.status = UploadStatus.CANCELLED
        self._active_tasks.clear()
        logger.info("HF Uploader shut down")

    async def _run_upload(
        self,
        task_id: str,
        token: str,
        readme_source_path: str,
        auto_readme: bool,
        redownload_notice: bool,
        private: bool,
    ) -> None:
        """Execute an upload task with semaphore-guarded sequential processing."""
        from huggingface_hub import HfApi

        task = self._tasks[task_id]
        tmp_readme: Optional[Path] = None

        try:
            async with self._upload_sem:
                if task_id in self._cancelled:
                    return

                task.status = UploadStatus.UPLOADING
                task.started_at = time.time()

                model_path = Path(task.model_path)
                api = HfApi()

                # Create repo (exist_ok handles already-existing repos)
                await asyncio.to_thread(
                    api.create_repo,
                    repo_id=task.repo_id,
                    token=token,
                    exist_ok=True,
                    private=private,
                )

                if task_id in self._cancelled:
                    return

                # Handle README
                readme_in_model = model_path / "README.md"
                if readme_source_path:
                    source_readme = Path(readme_source_path) / "README.md"
                    if source_readme.exists():
                        shutil.copy2(source_readme, readme_in_model)
                        tmp_readme = readme_in_model
                elif auto_readme and not _has_meaningful_readme(model_path):
                    try:
                        with open(model_path / "config.json") as f:
                            config = json.load(f)
                    except Exception:
                        config = {}
                    readme_content = _generate_model_card(
                        task.model_name, config,
                        redownload_notice=redownload_notice,
                    )
                    readme_in_model.write_text(readme_content, encoding="utf-8")
                    tmp_readme = readme_in_model

                if task_id in self._cancelled:
                    return

                # Upload the entire model folder
                # upload_folder is blocking; run in thread
                task.progress = 10.0  # Signal that upload has started

                await asyncio.to_thread(
                    api.upload_folder,
                    folder_path=str(model_path),
                    repo_id=task.repo_id,
                    token=token,
                    commit_message=f"Upload {task.model_name} via oMLX",
                )

                if task_id in self._cancelled:
                    return

                # Success
                task.status = UploadStatus.COMPLETED
                task.progress = 100.0
                task.completed_at = time.time()
                task.repo_url = f"https://huggingface.co/{task.repo_id}"

                elapsed = task.completed_at - task.started_at
                logger.info(
                    f"Upload completed: {task.model_name} -> {task.repo_id} "
                    f"({elapsed:.0f}s, {_format_size(task.total_size)})"
                )

        except asyncio.CancelledError:
            if task.status not in (UploadStatus.CANCELLED, UploadStatus.FAILED):
                task.status = UploadStatus.CANCELLED
        except Exception as e:
            if task_id not in self._cancelled:
                task.status = UploadStatus.FAILED
                task.error = str(e)
                logger.error(f"Upload failed for {task.model_name}: {e}")
        finally:
            # Clean up copied/generated README if we created it
            if tmp_readme and tmp_readme.exists():
                try:
                    tmp_readme.unlink()
                except Exception:
                    pass
            self._active_tasks.pop(task_id, None)


# ===== FILE: omlx/admin/ms_downloader.py =====

# SPDX-License-Identifier: Apache-2.0
"""ModelScope model downloader for oMLX admin panel.

Downloads models from ModelScope Hub using the modelscope SDK's snapshot_download
with directory-size-based progress polling.
"""

import asyncio
import logging
import os
import shutil
import time
import uuid
from pathlib import Path
from typing import Callable, Optional

import requests

from .hf_downloader import (
    DownloadStatus,
    DownloadTask,
    _format_model_size,
    _format_param_count,
)

logger = logging.getLogger(__name__)

# Check if modelscope SDK is available
MS_SDK_AVAILABLE = False
try:
    from modelscope import snapshot_download as ms_snapshot_download
    from modelscope.hub.api import HubApi as MSHubApi

    MS_SDK_AVAILABLE = True
except ImportError:
    ms_snapshot_download = None  # type: ignore[assignment]
    MSHubApi = None  # type: ignore[assignment, misc]

# Timeout for ModelScope API calls (seconds).
_MS_API_TIMEOUT = 15

# Seconds with no download progress before considering the download stalled.
_STALL_TIMEOUT = 300

# Default ModelScope API base URL.
_DEFAULT_MS_ENDPOINT = "https://modelscope.cn"

# Minimum downloads to be included in recommendations.
_MIN_DOWNLOADS = 50


def _get_ms_endpoint() -> str:
    """Get the configured ModelScope endpoint URL."""
    # Check environment variable first (set by CLI/settings)
    endpoint = os.environ.get("MODELSCOPE_DOMAIN", "")
    if endpoint:
        return endpoint.rstrip("/")

    try:
        from ..settings import get_settings

        endpoint = get_settings().modelscope.endpoint
        if endpoint:
            return endpoint.rstrip("/")
    except (RuntimeError, AttributeError):
        pass

    return _DEFAULT_MS_ENDPOINT


def _get_ms_api():
    """Create ModelScope HubApi instance.

    Returns:
        HubApi instance or None if SDK not available.
    """
    if not MS_SDK_AVAILABLE:
        return None
    return MSHubApi()


def _extract_model_size_from_files(file_list: list) -> int:
    """Calculate total file size from a list of file metadata dicts."""
    total = 0
    for f in file_list:
        size = f.get("Size") or f.get("size") or 0
        if isinstance(size, (int, float)):
            total += int(size)
    return total


# ---------------------------------------------------------------------------
# Per-model enrichment (size + param count)
#
# ModelScope's list_models endpoint returns Path/Name/Downloads/Likes/Stars
# but rarely populates StorageSize, and never returns a parameter count.
# To match HuggingFace's recommended-models card data, we enrich each entry
# with a config.json fetch (for params) and — when StorageSize was missing —
# a model-detail fetch (for size).
#
# Cached in-process for 24 hours since config.json content for a model
# doesn't change in practice; this keeps subsequent page loads of the
# Downloads tab essentially free.

_ENRICH_CACHE: dict[str, tuple[float, dict]] = {}
_ENRICH_CACHE_TTL = 24 * 3600  # 24h — config.json is effectively immutable
_ENRICH_CACHE_MAX = 1024       # bound memory under aggressive search/list use
_ENRICH_CONCURRENCY = 8        # parallel fetches per recommended/search call


def _enrich_cache_get(model_id: str) -> Optional[dict]:
    entry = _ENRICH_CACHE.get(model_id)
    if entry is None:
        return None
    ts, data = entry
    if time.time() - ts > _ENRICH_CACHE_TTL:
        _ENRICH_CACHE.pop(model_id, None)
        return None
    return data


def _enrich_cache_put(model_id: str, data: dict) -> None:
    if len(_ENRICH_CACHE) >= _ENRICH_CACHE_MAX:
        # Drop the oldest entry. O(N) on eviction but N is bounded at MAX
        # and evictions are rare in practice (24h TTL >> page-load rate).
        oldest = min(_ENRICH_CACHE, key=lambda k: _ENRICH_CACHE[k][0])
        _ENRICH_CACHE.pop(oldest, None)
    _ENRICH_CACHE[model_id] = (time.time(), data)


def _estimate_params_from_config(config: Optional[dict]) -> int:
    """Estimate decoder-transformer parameter count from a HF-style config.

    Handles dense Llama/Qwen/Mistral families and MoE variants
    (num_local_experts / num_experts). Returns 0 when required fields are
    missing — caller should render a blank rather than display a wrong
    number. The estimate is intentionally a rough headline figure (≈±5%);
    the goal is to surface "~7B" vs "~14B", not to match the checkpoint
    byte-for-byte.
    """
    if not isinstance(config, dict):
        return 0
    try:
        vocab_size = int(config.get("vocab_size", 0))
        hidden_size = int(config.get("hidden_size", 0))
        num_layers = int(config.get("num_hidden_layers", 0))
    except (TypeError, ValueError):
        return 0

    if not (vocab_size and hidden_size and num_layers):
        return 0

    try:
        intermediate_size = int(config.get("intermediate_size", 0))
        num_heads = int(config.get("num_attention_heads", 0))
        num_kv = int(config.get("num_key_value_heads", num_heads))
        head_dim = int(config.get("head_dim", 0)) or (
            hidden_size // num_heads if num_heads else 0
        )
        num_experts = int(
            config.get("num_local_experts")
            or config.get("num_experts")
            or 1
        )
        tie_embeddings = bool(config.get("tie_word_embeddings", True))
    except (TypeError, ValueError):
        return 0

    embeddings = vocab_size * hidden_size

    # Attention: Q + O are full hidden_size; K + V are reduced for GQA.
    if num_heads and head_dim:
        attn = (
            2 * hidden_size * (num_heads * head_dim)
            + 2 * hidden_size * (num_kv * head_dim)
        )
    else:
        attn = 4 * hidden_size * hidden_size

    # Gated MLP (Llama/Qwen style): gate + up + down projections.
    # MoE multiplies the FFN by the number of experts.
    if intermediate_size:
        ffn = num_experts * 3 * hidden_size * intermediate_size
    else:
        ffn = 8 * hidden_size * hidden_size

    layer_norms = 2 * hidden_size
    per_layer = attn + ffn + layer_norms

    total = embeddings + num_layers * per_layer + hidden_size
    if not tie_embeddings:
        total += vocab_size * hidden_size  # untied LM head

    return total


async def _fetch_model_config(model_id: str) -> Optional[dict]:
    """Fetch and parse a model's config.json from ModelScope.

    Returns None on any error (network, non-200, non-JSON) so callers can
    treat the field as absent without raising.
    """
    if not model_id:
        return None
    import json

    endpoint = _get_ms_endpoint()
    url = (
        f"{endpoint}/api/v1/models/{model_id}/repo"
        f"?FilePath=config.json&Revision=master"
    )
    try:
        resp = await asyncio.wait_for(
            asyncio.to_thread(requests.get, url, timeout=_MS_API_TIMEOUT),
            timeout=_MS_API_TIMEOUT + 5,
        )
        if resp.status_code != 200:
            return None
        return json.loads(resp.text)
    except Exception as e:
        logger.debug(f"config.json fetch failed for {model_id}: {e}")
        return None


async def _fetch_model_detail_size(model_id: str) -> int:
    """Fetch a model's storage size via the detail endpoint.

    Used as a fallback when list_models didn't populate StorageSize.
    Prefers ModelInfos.safetensor.model_size (weights only) and falls
    back to the repository StorageSize (weights + tokenizer + readme).
    Returns 0 on any error.
    """
    if not model_id:
        return 0
    endpoint = _get_ms_endpoint()
    url = f"{endpoint}/api/v1/models/{model_id}"
    try:
        resp = await asyncio.wait_for(
            asyncio.to_thread(requests.get, url, timeout=_MS_API_TIMEOUT),
            timeout=_MS_API_TIMEOUT + 5,
        )
        if resp.status_code != 200:
            return 0
        data = resp.json().get("Data") or {}
        model_infos = data.get("ModelInfos") or {}
        st = model_infos.get("safetensor") or {}
        size = st.get("model_size") or data.get("StorageSize") or 0
        return int(size) if isinstance(size, (int, float, str)) else 0
    except (TypeError, ValueError):
        return 0
    except Exception as e:
        logger.debug(f"model detail fetch failed for {model_id}: {e}")
        return 0


async def _enrich_ms_entry(entry: dict, sem: asyncio.Semaphore) -> dict:
    """Add size + params to a parsed model entry.

    Concurrent fetches are gated by `sem`; per-model results are cached
    in-process for 24h so subsequent page loads don't re-issue requests.
    Mutates and returns the same dict for ergonomic gather() pipelines.
    """
    model_id = entry.get("repo_id") or ""
    if not model_id:
        return entry

    cached = _enrich_cache_get(model_id)
    if cached is not None:
        c_size = cached.get("size") or 0
        c_params = cached.get("params") or 0
        if c_size and not entry.get("size"):
            entry["size"] = c_size
            entry["size_formatted"] = _format_model_size(c_size)
        if c_params:
            entry["params"] = c_params
            entry["params_formatted"] = _format_param_count(c_params)
        return entry

    async with sem:
        config_task = asyncio.create_task(_fetch_model_config(model_id))
        need_size = (entry.get("size") or 0) <= 0
        detail_task = (
            asyncio.create_task(_fetch_model_detail_size(model_id))
            if need_size else None
        )

        config = await config_task
        params = _estimate_params_from_config(config)

        size = entry.get("size") or 0
        if detail_task is not None:
            size = await detail_task

    _enrich_cache_put(model_id, {"size": size, "params": params})

    if size and not entry.get("size"):
        entry["size"] = size
        entry["size_formatted"] = _format_model_size(size)
    if params:
        entry["params"] = params
        entry["params_formatted"] = _format_param_count(params)
    return entry


def _parse_ms_model_entry(entry: dict) -> dict:
    """Parse a ModelScope API model entry into a normalized dict.

    Args:
        entry: Raw model dict from ModelScope API.

    Returns:
        Normalized model dict matching the HF format.
    """
    # Path is the organization/owner, Name is the model name
    # repo_id should be "owner/model" format
    path = entry.get("Path") or ""
    name = entry.get("Name") or ""
    if path and name:
        model_id = f"{path}/{name}"
    elif name:
        model_id = name
    else:
        model_id = path

    downloads = entry.get("Downloads") or 0
    likes = entry.get("Likes") or entry.get("Stars") or 0
    # StorageSize is the total size in bytes
    size = entry.get("StorageSize") or 0

    return {
        "repo_id": model_id,
        "name": name or model_id.split("/")[-1],
        "downloads": downloads,
        "likes": likes,
        "trending_score": 0,
        "size": size,
        "size_formatted": _format_model_size(size) if size > 0 else "",
        "params": None,
        "params_formatted": None,
    }


async def _fetch_ms_models_rest(
    query: str = "",
    page_size: int = 200,
) -> list[dict]:
    """Fetch models from ModelScope REST API without org restriction.

    Used when mlx_only is disabled to search across all organizations.

    Args:
        query: Optional search query to filter by model name.
        page_size: Number of models to fetch.

    Returns:
        List of raw model entry dicts from the API response.
    """
    endpoint = _get_ms_endpoint()
    url = f"{endpoint}/api/v1/models/"
    payload: dict = {"PageSize": page_size}
    if query:
        payload["Name"] = query
    try:
        resp = await asyncio.wait_for(
            asyncio.to_thread(
                requests.put, url, json=payload, timeout=_MS_API_TIMEOUT
            ),
            timeout=_MS_API_TIMEOUT + 5,
        )
        if resp.status_code == 200:
            data = resp.json().get("Data", {})
            return data.get("Models", data.get("models", []))
    except Exception as e:
        logger.warning(f"ModelScope REST API fetch failed: {e}")
    return []


class MSDownloader:
    """Manages ModelScope model downloads with progress tracking.

    Uses modelscope.snapshot_download() for actual downloads and polls
    the target directory size to estimate progress.

    Args:
        model_dir: Directory where downloaded models are stored.
        on_complete: Async callback invoked when a download completes successfully.
    """

    @staticmethod
    async def get_recommended_models(
        max_memory_bytes: int,
        limit: int = 60,
        result_limit: int = 50,
        mlx_only: bool = True,
    ) -> dict:
        """Fetch trending and popular models from ModelScope.

        When mlx_only is True, uses SDK to list models from mlx-community
        organization. When False, uses REST API to search all organizations.

        Args:
            max_memory_bytes: Maximum model size in bytes (typically system memory).
            limit: Number of models to fetch per category.
            result_limit: Maximum number of models to return per category.
            mlx_only: If True, restrict to mlx-community organization.

        Returns:
            Dict with 'trending' and 'popular' lists.
        """

        async def _fetch() -> list[dict]:
            if mlx_only:
                api = _get_ms_api()
                if api is None:
                    logger.warning("ModelScope SDK not available")
                    return []
                try:
                    data = await asyncio.wait_for(
                        asyncio.to_thread(
                            api.list_models,
                            "mlx-community",
                            page_size=limit,
                        ),
                        timeout=_MS_API_TIMEOUT + 5,
                    )
                except Exception as e:
                    logger.warning(f"ModelScope recommended fetch failed: {e}")
                    return []
                models_data = data.get("Models", [])
                if not models_data:
                    models_data = data.get("models", [])
            else:
                models_data = await _fetch_ms_models_rest(page_size=limit)

            results = []
            for entry in models_data:
                m = _parse_ms_model_entry(entry)
                downloads = m.get("downloads", 0)
                size = m.get("size", 0)
                # Filter by minimum downloads
                if downloads < _MIN_DOWNLOADS:
                    continue
                # Filter by memory size (only when list_models already had
                # a size — enrichment may reveal more below).
                if size > 0 and size > max_memory_bytes:
                    continue
                results.append(m)
                if len(results) >= result_limit * 2:
                    break

            return results

        models = await _fetch()

        # Enrich with size + params from per-model config.json / detail
        # fetches. Bounded concurrency keeps the call to ~1–2s for a full
        # page; results are cached in-process so subsequent loads are free.
        if models:
            sem = asyncio.Semaphore(_ENRICH_CONCURRENCY)
            enriched = await asyncio.gather(
                *(_enrich_ms_entry(m, sem) for m in models),
                return_exceptions=True,
            )
            models = [m for m in enriched if isinstance(m, dict)]

            # Re-apply the memory filter now that enrichment may have
            # supplied a real size for entries that list_models reported
            # as 0. Entries that still have no size are kept (better to
            # show with a blank size than hide a candidate the user has
            # enough RAM for).
            models = [
                m for m in models
                if (m.get("size", 0) == 0) or (m["size"] <= max_memory_bytes)
            ]

        # Sort by downloads for popular, keep original order for trending
        trending = models[:result_limit]
        popular = sorted(models, key=lambda x: x.get("downloads", 0), reverse=True)[:result_limit]

        return {
            "trending": trending,
            "popular": popular,
        }

    @staticmethod
    async def search_models(
        query: str,
        sort: str = "trending",
        limit: int = 100,
        mlx_only: bool = True,
    ) -> dict:
        """Search models on ModelScope.

        When mlx_only is True, uses SDK to list models from mlx-community
        and filters by query string. When False, uses REST API to search
        across all organizations.

        Args:
            query: Search query string.
            sort: Sort order (trending/downloads/created/updated).
            limit: Maximum number of results to return.
            mlx_only: If True, restrict to mlx-community organization.

        Returns:
            Dict with 'models' list and 'total' count.
        """
        if mlx_only:
            api = _get_ms_api()
            if api is None:
                logger.warning("ModelScope SDK not available")
                return {"models": [], "total": 0}

            try:
                data = await asyncio.wait_for(
                    asyncio.to_thread(
                        api.list_models,
                        "mlx-community",
                        page_size=200,
                    ),
                    timeout=_MS_API_TIMEOUT + 5,
                )
            except Exception as e:
                logger.error(f"ModelScope search failed: {e}")
                return {"models": [], "total": 0}

            models_data = data.get("Models", [])
            if not models_data:
                models_data = data.get("models", [])
        else:
            models_data = await _fetch_ms_models_rest(
                query=query, page_size=200
            )

        # Filter by query string (case-insensitive)
        query_lower = query.lower()
        filtered = []
        for entry in models_data:
            name = entry.get("Name", "")
            if query_lower in name.lower():
                m = _parse_ms_model_entry(entry)
                filtered.append(m)

        # Sort results
        if sort == "downloads":
            filtered.sort(key=lambda x: x.get("downloads", 0), reverse=True)
        elif sort == "created":
            pass  # Keep original order (newest first by default)
        elif sort == "updated":
            pass  # Keep original order

        # Limit results
        results = filtered[:limit]

        return {
            "models": results,
            "total": len(filtered),
        }

    @staticmethod
    async def get_model_info(model_id: str) -> dict:
        """Fetch detailed model information from ModelScope.

        Args:
            model_id: ModelScope model ID (e.g., "qwen/Qwen2.5-7B-Instruct-MLX").

        Returns:
            Dict with model details including description, files, tags, etc.
        """
        api = _get_ms_api()
        if api is None:
            raise RuntimeError("ModelScope SDK not available")

        # Get model metadata
        model_data = await asyncio.wait_for(
            asyncio.to_thread(api.get_model, model_id),
            timeout=_MS_API_TIMEOUT,
        )

        # get_model may return a string or dict depending on SDK version
        if isinstance(model_data, str):
            import json

            try:
                model_data = json.loads(model_data)
            except (json.JSONDecodeError, TypeError):
                model_data = {}

        if not isinstance(model_data, dict):
            model_data = {}

        # Get file list
        files = []
        total_file_size = 0
        try:
            file_list = await asyncio.wait_for(
                asyncio.to_thread(api.get_model_files, model_id),
                timeout=_MS_API_TIMEOUT,
            )
            for f in file_list or []:
                fname = f.get("Name") or f.get("Path", "")
                fsize = f.get("Size") or 0
                if isinstance(fsize, str):
                    try:
                        fsize = int(fsize)
                    except ValueError:
                        fsize = 0
                total_file_size += fsize
                files.append(
                    {
                        "name": fname,
                        "size": fsize,
                        "size_formatted": (
                            _format_model_size(fsize) if fsize > 0 else ""
                        ),
                    }
                )
        except Exception as e:
            logger.warning(f"Could not fetch file list for {model_id}: {e}")

        # Fetch model card (README.md) content
        model_card = ""
        try:
            endpoint = _get_ms_endpoint()
            readme_url = f"{endpoint}/api/v1/models/{model_id}/repo?FilePath=README.md&Revision=master"
            resp = await asyncio.wait_for(
                asyncio.to_thread(
                    requests.get,
                    readme_url,
                    timeout=_MS_API_TIMEOUT,
                ),
                timeout=_MS_API_TIMEOUT + 5,
            )
            if resp.status_code == 200:
                card_text = resp.text
                # Strip YAML front matter (between --- markers)
                if card_text.startswith("---"):
                    end = card_text.find("---", 3)
                    if end != -1:
                        card_text = card_text[end + 3:].strip()
                model_card = card_text
        except Exception:
            pass  # README not available

        # Extract metadata
        name = model_data.get("Name") or model_id
        downloads = model_data.get("Downloads") or 0
        likes = model_data.get("Likes") or model_data.get("Stars") or 0
        tags = model_data.get("Tags") or []
        if isinstance(tags, str):
            tags = [t.strip() for t in tags.split(",") if t.strip()]

        return {
            "repo_id": model_id,
            "name": name,
            "model_card": model_card,
            "description": model_data.get("Description", ""),
            "files": files,
            "tags": tags,
            "pipeline_tag": model_data.get("Task", ""),
            "params": None,
            "params_formatted": None,
            "size": total_file_size,
            "size_formatted": (
                _format_model_size(total_file_size) if total_file_size > 0 else ""
            ),
            "downloads": downloads,
            "likes": likes,
            "created_at": model_data.get("CreatedTime", ""),
            "updated_at": model_data.get("LastUpdatedTime", ""),
        }

    def __init__(
        self,
        model_dir: str,
        on_complete: Optional[Callable] = None,
    ):
        self._model_dir = Path(model_dir)
        self._tasks: dict[str, DownloadTask] = {}
        self._active_tasks: dict[str, asyncio.Task] = {}
        self._progress_tasks: dict[str, asyncio.Task] = {}
        self._on_complete = on_complete
        self._cancelled: set[str] = set()
        self._download_sem = asyncio.Semaphore(1)

    @property
    def model_dir(self) -> Path:
        return self._model_dir

    def update_model_dir(self, new_dir: str) -> None:
        """Update the model directory path."""
        self._model_dir = Path(new_dir)

    async def start_download(
        self, model_id: str, ms_token: str = ""
    ) -> DownloadTask:
        """Start downloading a model from ModelScope.

        Args:
            model_id: ModelScope model ID (e.g., "qwen/Qwen2.5-7B-Instruct-MLX").
            ms_token: Optional ModelScope token for private models.

        Returns:
            The created DownloadTask.

        Raises:
            ValueError: If model_id format is invalid or download is already queued.
            RuntimeError: If ModelScope SDK is not installed.
        """
        if not MS_SDK_AVAILABLE:
            raise RuntimeError(
                "ModelScope SDK not installed. "
                "Install with: pip install \"omlx[modelscope]\""
            )

        model_id = model_id.strip()
        if "/" not in model_id or len(model_id.split("/")) != 2:
            raise ValueError(
                f"Invalid model ID: '{model_id}'. "
                "Expected format: 'owner/model' (e.g., 'qwen/Qwen2.5-7B-Instruct-MLX')"
            )

        # Check for duplicate active downloads
        for task in self._tasks.values():
            if task.repo_id == model_id and task.status in (
                DownloadStatus.PENDING,
                DownloadStatus.DOWNLOADING,
            ):
                raise ValueError(
                    f"Download for '{model_id}' is already in progress"
                )

        task_id = str(uuid.uuid4())
        task = DownloadTask(task_id=task_id, repo_id=model_id)
        self._tasks[task_id] = task

        # Start download in background
        self._active_tasks[task_id] = asyncio.create_task(
            self._run_download(task_id, ms_token)
        )

        logger.info(f"MS Download queued: {model_id} (task_id={task_id})")
        return task

    async def cancel_download(self, task_id: str) -> bool:
        """Cancel an active download.

        Note: Due to Python threading limitations, the actual download thread
        cannot be interrupted immediately. The download will be marked as
        cancelled and files will be cleaned up when the thread completes.

        Args:
            task_id: The task ID to cancel.

        Returns:
            True if the task was found and cancelled.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status not in (DownloadStatus.PENDING, DownloadStatus.DOWNLOADING):
            return False

        # Mark as cancelled - the running thread will check this flag
        self._cancelled.add(task_id)
        task.status = DownloadStatus.CANCELLED
        task.error = "Cancellation requested. Download will stop shortly."

        # Stop progress polling
        progress_task = self._progress_tasks.pop(task_id, None)
        if progress_task and not progress_task.done():
            progress_task.cancel()

        # Cancel the download task
        active_task = self._active_tasks.pop(task_id, None)
        if active_task and not active_task.done():
            active_task.cancel()

        logger.info(f"MS Download cancelled: {task.repo_id} (task_id={task_id})")
        return True

    def remove_task(self, task_id: str) -> bool:
        """Remove a completed, failed, or cancelled task from the list.

        Args:
            task_id: The task ID to remove.

        Returns:
            True if the task was found and removed.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return False

        if task.status in (DownloadStatus.PENDING, DownloadStatus.DOWNLOADING):
            return False

        del self._tasks[task_id]
        self._cancelled.discard(task_id)
        return True

    async def retry_download(
        self, task_id: str, ms_token: str = ""
    ) -> DownloadTask:
        """Retry a failed or cancelled download, resuming from existing files.

        Args:
            task_id: The task ID of the failed/cancelled download.
            ms_token: Optional ModelScope token for private models.

        Returns:
            The new DownloadTask.

        Raises:
            ValueError: If task not found or not in retryable state.
        """
        old_task = self._tasks.get(task_id)
        if old_task is None:
            raise ValueError(f"Task not found: {task_id}")

        if old_task.status not in (DownloadStatus.FAILED, DownloadStatus.CANCELLED):
            raise ValueError(
                f"Task {task_id} is not retryable (status: {old_task.status.value})"
            )

        model_id = old_task.repo_id
        old_retry_count = old_task.retry_count

        # Remove old task entry
        del self._tasks[task_id]
        self._cancelled.discard(task_id)

        # Start fresh download (snapshot_download resumes from existing files)
        new_task = await self.start_download(model_id, ms_token)
        new_task.retry_count = old_retry_count + 1
        return new_task

    def get_tasks(self) -> list[dict]:
        """Return all tasks as serializable dicts, ordered by creation time."""
        return [
            task.to_dict()
            for task in sorted(self._tasks.values(), key=lambda t: t.created_at)
        ]

    async def shutdown(self) -> None:
        """Cancel all active downloads and clean up."""
        # Cancel all progress polling tasks
        for task_id, progress_task in list(self._progress_tasks.items()):
            if not progress_task.done():
                progress_task.cancel()
        self._progress_tasks.clear()

        # Cancel all active download tasks
        for task_id, active_task in list(self._active_tasks.items()):
            if not active_task.done():
                active_task.cancel()
                task = self._tasks.get(task_id)
                if task and task.status == DownloadStatus.DOWNLOADING:
                    task.status = DownloadStatus.CANCELLED
        self._active_tasks.clear()

        logger.info("MS Downloader shut down")

    async def _run_download(self, task_id: str, ms_token: str) -> None:
        """Execute a download task.

        Waits for the download semaphore (only one download runs at a time),
        then fetches file info for total size and runs snapshot_download in a
        thread while polling the target directory for progress updates.
        """
        task = self._tasks[task_id]

        try:
            async with self._download_sem:
                # Check if cancelled while waiting in queue
                if task_id in self._cancelled:
                    return

                task.status = DownloadStatus.DOWNLOADING
                task.started_at = time.time()

                # Preserve {owner}/{model} layout to match other tools
                # (LMStudio, huggingface-cli) and avoid duplicate downloads
                # when sharing a model directory.
                target_dir = self._model_dir / task.repo_id

                # Get total file size for progress estimation
                try:
                    api = _get_ms_api()
                    if api:
                        file_list = await asyncio.wait_for(
                            asyncio.to_thread(api.get_model_files, task.repo_id),
                            timeout=_MS_API_TIMEOUT,
                        )
                        if file_list:
                            task.total_size = _extract_model_size_from_files(
                                file_list
                            )
                except Exception as e:
                    logger.warning(
                        f"Could not fetch file info for {task.repo_id}: {e}. "
                        "Progress estimation will be unavailable."
                    )

                # Start progress polling
                self._progress_tasks[task_id] = asyncio.create_task(
                    self._poll_progress(task_id, target_dir)
                )

                # Build download kwargs
                dl_kwargs = {
                    "model_id": task.repo_id,
                    "local_dir": str(target_dir),
                }
                if ms_token:
                    dl_kwargs["token"] = ms_token

                # Run snapshot_download in a thread (blocking call)
                # Note: Thread cannot be interrupted, cancellation is checked after completion
                await asyncio.to_thread(
                    ms_snapshot_download,
                    **dl_kwargs,
                )

                # Check if cancelled while downloading - clean up downloaded files
                if task_id in self._cancelled:
                    logger.info(
                        f"MS Download was cancelled during execution: {task.repo_id}. "
                        "Cleaning up downloaded files..."
                    )
                    if target_dir.exists():
                        try:
                            shutil.rmtree(target_dir)
                            logger.info(f"Cleaned up cancelled download: {target_dir}")
                        except Exception as cleanup_err:
                            logger.warning(f"Failed to clean up {target_dir}: {cleanup_err}")
                    # Drop empty org folder left behind by the cancelled download.
                    parent = target_dir.parent
                    if (
                        parent != self._model_dir
                        and parent.exists()
                        and not any(parent.iterdir())
                    ):
                        try:
                            parent.rmdir()
                        except OSError as cleanup_err:
                            logger.debug(
                                f"Could not remove empty org folder {parent}: "
                                f"{cleanup_err}"
                            )
                    return

                # Success
                task.status = DownloadStatus.COMPLETED
                task.progress = 100.0
                task.downloaded_size = task.total_size or self._get_dir_size(
                    target_dir
                )
                task.completed_at = time.time()

                logger.info(
                    f"MS Download completed: {task.repo_id} -> {target_dir} "
                    f"({time.time() - task.started_at:.1f}s)"
                )

                # Trigger model pool refresh
                if self._on_complete:
                    try:
                        await self._on_complete()
                    except Exception as e:
                        logger.error(
                            f"Error in download completion callback: {e}"
                        )

        except asyncio.CancelledError:
            if task.status not in (
                DownloadStatus.CANCELLED,
                DownloadStatus.FAILED,
            ):
                task.status = DownloadStatus.CANCELLED
        except Exception as e:
            if task_id not in self._cancelled:
                task.status = DownloadStatus.FAILED
                # Provide user-friendly error messages
                err_msg = str(e)
                if "NotExistError" in type(e).__name__ or "404" in err_msg:
                    task.error = (
                        f"Model not found: {task.repo_id}. "
                        "Please check the model ID and try again."
                    )
                elif "401" in err_msg or "403" in err_msg:
                    task.error = (
                        f"Access denied for '{task.repo_id}'. "
                        "Please provide a valid ModelScope token."
                    )
                else:
                    task.error = err_msg
                logger.error(f"MS Download failed for {task.repo_id}: {e}")
        finally:
            # Stop progress polling
            progress_task = self._progress_tasks.pop(task_id, None)
            if progress_task and not progress_task.done():
                progress_task.cancel()

            # Remove from active tasks
            self._active_tasks.pop(task_id, None)

    async def _poll_progress(self, task_id: str, target_dir: Path) -> None:
        """Poll the target directory to estimate download progress.

        Uses both directory size and file modification times to detect
        activity.
        """
        task = self._tasks.get(task_id)
        if task is None:
            return

        last_size = 0
        last_activity_at = time.time()

        try:
            while task.status == DownloadStatus.DOWNLOADING:
                await asyncio.sleep(2)

                if task.status != DownloadStatus.DOWNLOADING:
                    break

                current_size = self._get_dir_size(target_dir)
                task.downloaded_size = current_size

                if task.total_size > 0:
                    # Cap at 99% until snapshot_download confirms completion
                    task.progress = min(
                        (current_size / task.total_size) * 100, 99.0
                    )

                # Activity detection: size change OR file mtime change
                if current_size != last_size:
                    last_size = current_size
                    last_activity_at = time.time()
                else:
                    latest_mtime = self._get_latest_mtime(target_dir)
                    if latest_mtime > last_activity_at:
                        last_activity_at = latest_mtime

                # Stall detection
                if (
                    current_size > 0
                    and (time.time() - last_activity_at) > _STALL_TIMEOUT
                ):
                    task.status = DownloadStatus.FAILED
                    task.error = (
                        f"Download stalled: no progress for {_STALL_TIMEOUT}s. "
                        "Try retrying the download."
                    )
                    logger.warning(
                        f"MS Download stalled for {task.repo_id} "
                        f"(task_id={task_id})"
                    )
                    # Cancel the snapshot_download thread
                    active_task = self._active_tasks.get(task_id)
                    if active_task and not active_task.done():
                        active_task.cancel()
                    break
        except asyncio.CancelledError:
            pass

    @staticmethod
    def _get_latest_mtime(path: Path) -> float:
        """Return the most recent modification time of any file in a directory."""
        if not path.exists():
            return 0.0
        latest = 0.0
        try:
            for f in path.rglob("*"):
                if f.is_file():
                    try:
                        mt = f.stat().st_mtime
                        if mt > latest:
                            latest = mt
                    except OSError:
                        pass
        except OSError:
            pass
        return latest

    @staticmethod
    def _get_dir_size(path: Path) -> int:
        """Calculate total size of all files in a directory."""
        if not path.exists():
            return 0
        total = 0
        try:
            for f in path.rglob("*"):
                if f.is_file():
                    try:
                        total += f.stat().st_size
                    except OSError:
                        pass
        except OSError:
            pass
        return total


# ===== FILE: omlx/admin/oq_manager.py =====

# SPDX-License-Identifier: Apache-2.0
"""oQ Quantization task manager for the admin panel.

Manages quantization tasks with progress tracking, following the same pattern
as hf_downloader.py (DownloadTask / HFDownloader).
"""

import asyncio
import enum
import hashlib
import json
import logging
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

try:
    import mlx.core as mx

    HAS_MLX = True
except ImportError:
    HAS_MLX = False

from ..model_discovery import _has_vision_subconfig

logger = logging.getLogger(__name__)


class _QuantCancelled(Exception):
    """Raised by progress callback when task is cancelled."""

    pass


class QuantStatus(str, enum.Enum):
    """Status of a quantization task."""

    PENDING = "pending"
    LOADING = "loading"
    QUANTIZING = "quantizing"
    SAVING = "saving"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


_ACTIVE_STATUSES = {
    QuantStatus.PENDING,
    QuantStatus.LOADING,
    QuantStatus.QUANTIZING,
    QuantStatus.SAVING,
}


@dataclass
class QuantTask:
    """Represents a single oQ quantization task."""

    task_id: str
    model_name: str
    model_path: str
    oq_level: float
    output_name: str
    output_path: str
    status: QuantStatus = QuantStatus.PENDING
    progress: float = 0.0
    phase: str = ""
    progress_detail: str = ""
    progress_meta: dict = field(default_factory=dict)
    error: str = ""
    created_at: float = field(default_factory=time.time)
    started_at: float = 0.0
    completed_at: float = 0.0
    source_size: int = 0
    output_size: int = 0
    group_size: int = 64
    sensitivity_model_path: str = ""
    text_only: bool = False
    dtype: str = "bfloat16"
    preserve_mtp: bool = False
    auto_proxy_sensitivity: bool = True
    enhanced: bool = False
    imatrix_cache_path: str = ""
    imatrix_reuse_cache: bool = True
    imatrix_strict: bool = False
    imatrix_num_samples: int = 128
    imatrix_seq_length: int = 512
    mtp_assistant_model_path: str = ""

    def to_dict(self) -> dict:
        """Serialize task to JSON-compatible dict."""
        return {
            "task_id": self.task_id,
            "model_name": self.model_name,
            "model_path": self.model_path,
            "oq_level": self.oq_level,
            "output_name": self.output_name,
            "output_path": self.output_path,
            "status": self.status.value,
            "progress": round(self.progress, 1),
            "phase": self.phase,
            "progress_detail": self.progress_detail,
            "progress_meta": self.progress_meta,
            "error": self.error,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "source_size": self.source_size,
            "output_size": self.output_size,
            "dtype": self.dtype,
            "enhanced": self.enhanced,
            "imatrix_cache_path": self.imatrix_cache_path,
        }


def _dir_size(path: Path) -> int:
    """Get total size of files in a directory."""
    if not path.exists():
        return 0
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())


def _format_size(size_bytes: int) -> str:
    """Format byte count as human-readable string."""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024**2:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024**3:
        return f"{size_bytes / 1024**2:.1f} MB"
    else:
        return f"{size_bytes / 1024**3:.1f} GB"


class OQManager:
    """Manages oQ quantization tasks with async execution and progress tracking.

    Follows the same pattern as HFDownloader: semaphore-guarded sequential
    execution, polling-based progress, cooperative cancellation.
    """

    def __init__(
        self,
        model_dirs: list[str],
        on_complete: Optional[Callable] = None,
    ):
        self._model_dirs = [Path(d) for d in model_dirs]
        self._output_dir = self._model_dirs[0] if self._model_dirs else Path(".")
        self._tasks: dict[str, QuantTask] = {}
        self._active_tasks: dict[str, asyncio.Task] = {}
        self._progress_tasks: dict[str, asyncio.Task] = {}
        self._on_complete = on_complete
        self._cancelled: set[str] = set()
        self._quant_sem = asyncio.Semaphore(1)

    def update_model_dirs(self, model_dirs: list[str]) -> None:
        """Update model directory paths."""
        self._model_dirs = [Path(d) for d in model_dirs]
        if self._model_dirs:
            self._output_dir = self._model_dirs[0]

    async def list_quantizable_models(self) -> tuple[list[dict], list[dict]]:
        """Scan all model dirs. Returns (source_models, all_models)."""

        def _scan() -> tuple[list[dict], list[dict]]:
            from ..oq import estimate_memory, validate_quantizable
            from ..utils.model_loading import (
                _checkpoint_has_mtp_weights,
                _has_mtp_heads,
            )

            source_models = []
            all_models = []
            seen: set[str] = set()

            for model_dir in self._model_dirs:
                if not model_dir.exists():
                    continue
                for subdir in sorted(model_dir.iterdir()):
                    if not subdir.is_dir():
                        continue
                    candidates = []
                    if (subdir / "config.json").exists():
                        candidates.append(subdir)
                    else:
                        for child in sorted(subdir.iterdir()):
                            if child.is_dir() and (child / "config.json").exists():
                                candidates.append(child)

                    for path in candidates:
                        if path.name in seen:
                            continue
                        seen.add(path.name)
                        try:
                            with open(path / "config.json") as f:
                                config = json.load(f)
                            size = sum(
                                f.stat().st_size for f in path.glob("*.safetensors")
                            )
                            if size == 0:
                                size = sum(f.stat().st_size for f in path.glob("*.bin"))
                            if size == 0:
                                continue
                            tc = config.get("text_config", {})
                            has_mtp = _has_mtp_heads(
                                config
                            ) and _checkpoint_has_mtp_weights(path)
                            info = {
                                "name": path.name,
                                "path": str(path),
                                "size": size,
                                "size_formatted": _format_size(size),
                                "model_type": config.get("model_type", "")
                                or tc.get("model_type", ""),
                                "is_quantized": "quantization" in config,
                                # Treat vision_config / vit_config / mm_vision_tower as VLM
                                # evidence (Molmo / Molmo2 use vit_config; FastVLM uses
                                # mm_vision_tower). Same predicate as model_discovery.
                                "is_vlm": _has_vision_subconfig(config),
                                "has_mtp_heads": has_mtp,
                                "hidden_size": tc.get("hidden_size")
                                or config.get("hidden_size")
                                or 0,
                            }
                            all_models.append(info)
                            if validate_quantizable(config):
                                info_full = dict(info)
                                info_full["num_layers"] = config.get(
                                    "num_hidden_layers", 0
                                ) or tc.get("num_hidden_layers", 0)
                                info_full["num_experts"] = config.get(
                                    "num_local_experts", 0
                                )
                                info_full["memory_streaming"] = estimate_memory(size)
                                source_models.append(info_full)
                        except Exception:
                            continue
            return source_models, all_models

        return await asyncio.to_thread(_scan)

    async def start_quantization(
        self,
        model_path: str,
        oq_level: float,
        group_size: int = 64,
        sensitivity_model_path: str = "",
        text_only: bool = False,
        dtype: str = "bfloat16",
        preserve_mtp: bool = False,
        auto_proxy_sensitivity: bool = True,
        enhanced: bool = False,
        imatrix_cache_path: str = "",
        imatrix_reuse_cache: bool = True,
        imatrix_strict: bool = False,
        imatrix_num_samples: int = 128,
        imatrix_seq_length: int = 512,
        mtp_assistant_model_path: str = "",
    ) -> QuantTask:
        """Start a quantization job.

        Args:
            model_path: Path to source model directory.
            oq_level: oQ level from OQ_LEVELS.
            dtype: Target fp dtype for non-quantized weights and quant
                scales/biases. "bfloat16" (default) or "float16".
            mtp_assistant_model_path: Optional checkpoint whose MTP head is
                merged into the output. A gemma4_assistant donor uses the
                assistant merge; any other donor grafts its native
                Qwen3.5/3.6 mtp.* head (same-geometry, same-tokenizer
                pairs only). Validated at submission.

        Returns:
            The created QuantTask.

        Raises:
            ValueError: On invalid inputs or output conflict.
        """
        from ..oq import (
            OQ_DTYPES,
            OQ_LEVELS,
            _validate_oq_dtype_for_model,
            resolve_output_name,
            validate_gemma4_assistant_pair,
            validate_mtp_donor_pair,
        )
        from ..utils.model_loading import _checkpoint_has_mtp_weights

        if oq_level not in OQ_LEVELS:
            raise ValueError(
                f"Invalid oQ level {oq_level}. Must be one of {sorted(OQ_LEVELS)}"
            )
        if dtype not in OQ_DTYPES:
            raise ValueError(f"Invalid dtype {dtype!r}. Must be one of {OQ_DTYPES}")

        source = Path(model_path)
        if not source.exists() or not (source / "config.json").exists():
            raise ValueError(f"Model not found: {model_path}")

        with open(source / "config.json") as f:
            config = json.load(f)
        _validate_oq_dtype_for_model(config, dtype)

        if preserve_mtp and not _checkpoint_has_mtp_weights(source):
            logger.warning(
                "Preserve MTP requested for %s, but no mtp.* tensors were "
                "found in the checkpoint; disabling MTP preservation",
                source.name,
            )
            preserve_mtp = False

        if preserve_mtp and mtp_assistant_model_path:
            raise ValueError(
                "Choose either 'Preserve MTP weights' or 'Combine MTP head', "
                "not both"
            )
        if mtp_assistant_model_path:
            assistant = Path(mtp_assistant_model_path)
            if not assistant.exists() or not (assistant / "config.json").exists():
                raise ValueError(
                    f"Assistant model not found: {mtp_assistant_model_path}"
                )
            with open(assistant / "config.json") as f:
                assistant_config = json.load(f)
            if assistant_config.get("model_type") == "gemma4_assistant":
                validate_gemma4_assistant_pair(config, assistant_config)
            else:
                validate_mtp_donor_pair(source, assistant)
                if _checkpoint_has_mtp_weights(source):
                    logger.warning(
                        "Recipient %s ships its own MTP head; it will be "
                        "stripped and replaced by the donor head from %s",
                        source.name,
                        assistant.name,
                    )

        model_name = source.name
        output_name = resolve_output_name(
            model_name,
            oq_level,
            dtype,
            preserve_mtp=preserve_mtp,
            enhanced=enhanced,
        )
        if mtp_assistant_model_path and not output_name.endswith("-mtp"):
            output_name += "-mtp"
        output_path = self._output_dir / output_name

        if output_path.exists():
            raise ValueError(
                f"Output directory already exists: {output_path}. "
                "Delete it first via the Manager tab."
            )

        # Check for duplicate active tasks (same level + dtype combo)
        for task in self._tasks.values():
            if (
                task.model_path == model_path
                and task.oq_level == oq_level
                and task.dtype == dtype
                and task.enhanced == enhanced
                and task.status in _ACTIVE_STATUSES
            ):
                raise ValueError(
                    f"Quantization for '{model_name}' at oQ{oq_level:g}"
                    f"{'e' if enhanced else ''} "
                    f"({dtype}) is already in progress"
                )

        if enhanced:
            if imatrix_num_samples < 1:
                raise ValueError("imatrix_num_samples must be >= 1")
            if imatrix_seq_length < 1:
                raise ValueError("imatrix_seq_length must be >= 1")
            if not imatrix_cache_path:
                digest = hashlib.sha256(str(source.resolve()).encode()).hexdigest()[:12]
                imatrix_cache_path = str(
                    self._output_dir
                    / ".oqe_imatrix"
                    / (
                        f"{model_name}-{digest}-s{int(imatrix_num_samples)}"
                        f"-l{int(imatrix_seq_length)}.npz"
                    )
                )

        source_size = sum(f.stat().st_size for f in source.glob("*.safetensors"))
        if source_size == 0:
            source_size = sum(f.stat().st_size for f in source.glob("*.bin"))

        task_id = str(uuid.uuid4())
        task = QuantTask(
            task_id=task_id,
            model_name=model_name,
            model_path=model_path,
            oq_level=oq_level,
            output_name=output_name,
            output_path=str(output_path),
            source_size=source_size,
            group_size=group_size,
            sensitivity_model_path=sensitivity_model_path,
            text_only=text_only,
            dtype=dtype,
            preserve_mtp=preserve_mtp,
            auto_proxy_sensitivity=auto_proxy_sensitivity,
            enhanced=enhanced,
            imatrix_cache_path=imatrix_cache_path,
            imatrix_reuse_cache=imatrix_reuse_cache,
            imatrix_strict=imatrix_strict,
            imatrix_num_samples=imatrix_num_samples,
            imatrix_seq_length=imatrix_seq_length,
            mtp_assistant_model_path=mtp_assistant_model_path,
        )
        self._tasks[task_id] = task

        self._active_tasks[task_id] = asyncio.create_task(
            self._run_quantization(task_id)
        )

        logger.info(
            f"oQ quantization queued: {model_name} -> "
            f"oQ{oq_level:g}{'e' if enhanced else ''} "
            f"(task_id={task_id})"
        )
        return task

    async def cancel_quantization(self, task_id: str) -> bool:
        """Cancel an active quantization task."""
        task = self._tasks.get(task_id)
        if task is None:
            return False
        if task.status not in _ACTIVE_STATUSES:
            return False

        self._cancelled.add(task_id)
        task.status = QuantStatus.CANCELLED

        progress_task = self._progress_tasks.pop(task_id, None)
        if progress_task and not progress_task.done():
            progress_task.cancel()

        active_task = self._active_tasks.pop(task_id, None)

        # Clean up partial output
        output = Path(task.output_path)
        if output.exists():
            import shutil

            shutil.rmtree(output, ignore_errors=True)

        # Wait for the quantization thread to actually finish.
        # Do NOT call active_task.cancel() first — that only cancels the
        # asyncio wrapper and causes the await to return immediately while
        # the OS thread continues running Metal commands. Instead, rely on
        # cooperative cancellation: the progress callback raises
        # _QuantCancelled when it sees the flag, terminating quantize_oq
        # at the next callback point (per-layer in GPTQ, per-tensor in
        # streaming).
        if active_task and not active_task.done():
            try:
                await asyncio.wait_for(
                    asyncio.shield(active_task),
                    timeout=30.0,
                )
            except asyncio.TimeoutError:
                # Thread didn't exit cooperatively (e.g. stuck in long GPTQ
                # block). Force-cancel as last resort and wait a bit for
                # Metal to settle.
                logger.warning(
                    "oQ cancel: cooperative exit timed out, force-cancelling"
                )
                active_task.cancel()
                try:
                    await active_task
                except (asyncio.CancelledError, Exception):
                    pass
                await asyncio.sleep(2.0)
            except (asyncio.CancelledError, Exception):
                pass

        # GPU cleanup after thread is done
        if HAS_MLX:
            for _attempt in range(3):
                try:
                    mx.synchronize()
                    mx.clear_cache()
                    break
                except Exception:
                    await asyncio.sleep(1.0)

        logger.info(f"oQ quantization cancelled: {task.model_name} (task_id={task_id})")
        return True

    def remove_task(self, task_id: str) -> bool:
        """Remove a completed/failed/cancelled task from the list."""
        task = self._tasks.get(task_id)
        if task is None:
            return False
        if task.status in _ACTIVE_STATUSES:
            return False
        del self._tasks[task_id]
        self._cancelled.discard(task_id)
        return True

    def get_tasks(self) -> list[dict]:
        """Return all tasks as serializable dicts."""
        return [t.to_dict() for t in self._tasks.values()]

    @property
    def is_quantizing(self) -> bool:
        """Check if any quantization task is actively running."""
        return any(t.status in _ACTIVE_STATUSES for t in self._tasks.values())

    async def shutdown(self) -> None:
        """Cancel all active tasks."""
        for task_id in list(self._active_tasks):
            await self.cancel_quantization(task_id)

    async def _run_quantization(self, task_id: str) -> None:
        """Execute the quantization pipeline in background."""
        task = self._tasks[task_id]
        try:
            async with self._quant_sem:
                if task_id in self._cancelled:
                    return

                # Ensure GPU is clean before starting (previous task may have been cancelled)
                # Metal command buffers need full sync + cache clear after cancellation
                if HAS_MLX:
                    for _ in range(3):
                        try:
                            mx.synchronize()
                            mx.clear_cache()
                            break
                        except Exception:
                            await asyncio.sleep(1.0)

                # Phase 1: Loading
                task.status = QuantStatus.LOADING
                task.started_at = time.time()
                task.phase = "Loading model..."
                task.progress = 5.0

                def _progress_cb(
                    phase: str,
                    pct: float,
                    detail: str = "",
                    meta: dict | None = None,
                ) -> None:
                    if task_id in self._cancelled:
                        raise _QuantCancelled(f"Task {task_id} cancelled")
                    base_phase = phase.split("|", 1)[0]
                    if base_phase.startswith("quantizing"):
                        task.status = QuantStatus.QUANTIZING
                    elif base_phase == "saving":
                        task.status = QuantStatus.SAVING
                    else:
                        task.status = QuantStatus.LOADING
                    task.phase = self._phase_label(phase, task.oq_level, task.enhanced)
                    task.progress_detail = detail or ""
                    task.progress_meta = meta or {}
                    task.progress = pct
                    task._last_progress_callback_at = time.time()

                # Start time-based progress estimation
                self._progress_tasks[task_id] = asyncio.create_task(
                    self._estimate_progress(task_id)
                )

                from ..oq import quantize_oq_streaming

                await asyncio.to_thread(
                    quantize_oq_streaming,
                    task.model_path,
                    task.output_path,
                    task.oq_level,
                    task.group_size,
                    _progress_cb,
                    task.text_only,
                    None,  # target_bpw
                    None,  # hard_cap_bpw
                    task.sensitivity_model_path,
                    task.dtype,
                    task.preserve_mtp,
                    task.auto_proxy_sensitivity,
                    enhanced=task.enhanced,
                    imatrix_cache_path=task.imatrix_cache_path,
                    imatrix_reuse_cache=task.imatrix_reuse_cache,
                    imatrix_strict=task.imatrix_strict,
                    imatrix_num_samples=task.imatrix_num_samples,
                    imatrix_seq_length=task.imatrix_seq_length,
                )

                if task_id in self._cancelled:
                    return

                if task.mtp_assistant_model_path:
                    from ..oq import combine_mtp_into_output

                    _progress_cb("saving", 97.0, "Merging MTP head...")
                    await asyncio.to_thread(
                        combine_mtp_into_output,
                        task.output_path,
                        task.mtp_assistant_model_path,
                    )

                if task_id in self._cancelled:
                    return

                # Complete
                task.status = QuantStatus.COMPLETED
                task.progress = 100.0
                task.phase = "Completed"
                task.progress_detail = ""
                task.progress_meta = {}
                task.completed_at = time.time()
                task.output_size = _dir_size(Path(task.output_path))

                elapsed = task.completed_at - task.started_at
                logger.info(
                    f"oQ quantization completed: {task.output_name} "
                    f"({elapsed:.0f}s, {_format_size(task.output_size)})"
                )

                if self._on_complete:
                    try:
                        result = self._on_complete()
                        if asyncio.iscoroutine(result):
                            await result
                    except Exception:
                        logger.exception("on_complete callback failed")

        except asyncio.CancelledError:
            if task.status not in (QuantStatus.CANCELLED, QuantStatus.FAILED):
                task.status = QuantStatus.CANCELLED
        except _QuantCancelled:
            if task.status != QuantStatus.CANCELLED:
                task.status = QuantStatus.CANCELLED
        except Exception as e:
            if task_id not in self._cancelled:
                task.status = QuantStatus.FAILED
                task.error = str(e)
                task.completed_at = time.time()
                logger.exception(f"oQ quantization failed: {task.model_name} -> {e}")
                # Clean up partial output
                output = Path(task.output_path)
                if output.exists():
                    import shutil

                    shutil.rmtree(output, ignore_errors=True)
        finally:
            pt = self._progress_tasks.pop(task_id, None)
            if pt and not pt.done():
                pt.cancel()
            self._active_tasks.pop(task_id, None)

    async def _estimate_progress(self, task_id: str) -> None:
        """Estimate progress by time during quantize phase (30-90%)."""
        task = self._tasks.get(task_id)
        if task is None:
            return

        source_gb = max(task.source_size / (1024**3), 0.1)
        estimated_total = source_gb * 3.0
        start = time.time()

        try:
            while task_id not in self._cancelled and task.status in _ACTIVE_STATUSES:
                await asyncio.sleep(2)
                elapsed = time.time() - start
                if time.time() - getattr(task, "_last_progress_callback_at", 0.0) < 5:
                    continue
                if task.status == QuantStatus.QUANTIZING:
                    if self._has_explicit_quant_progress(task):
                        continue
                    fraction = min(elapsed / estimated_total, 0.95)
                    task.progress = max(task.progress, 30.0 + fraction * 60.0)
                elif task.status == QuantStatus.SAVING:
                    # During save, poll output dir size
                    output = Path(task.output_path)
                    if output.exists() and task.source_size > 0:
                        current = _dir_size(output)
                        # Estimate output as source * (oq_level / 16)
                        expected = task.source_size * task.oq_level / 16
                        if expected > 0:
                            save_frac = min(current / expected, 0.99)
                            task.progress = max(task.progress, 90.0 + save_frac * 10.0)
        except asyncio.CancelledError:
            pass

    @staticmethod
    def _has_explicit_quant_progress(task: QuantTask) -> bool:
        """Return True once the quantizer emits byte-level progress."""
        meta = task.progress_meta if isinstance(task.progress_meta, dict) else {}
        try:
            total_bytes = int(meta.get("total_bytes") or 0)
            processed_bytes = int(meta.get("processed_bytes") or 0)
        except (TypeError, ValueError):
            return False
        return total_bytes > 0 and processed_bytes >= 0

    @staticmethod
    def _phase_label(phase: str, oq_level: float, enhanced: bool = False) -> str:
        """Human-readable phase label."""
        oq_label = f"oQ{oq_level:g}{'e' if enhanced else ''}"
        labels = {
            "loading": "Loading model...",
            "imatrix": "Collecting oQe imatrix...",
            "quantizing": f"Quantizing to {oq_label}...",
            "saving": "Saving quantized model...",
        }
        # Handle progress: "quantizing_eta|792|879|0:02"
        if phase.startswith("quantizing_eta|"):
            parts = phase.split("|")
            current = parts[1] if len(parts) > 1 else "?"
            total = parts[2] if len(parts) > 2 else "?"
            eta = parts[3] if len(parts) > 3 and parts[3] else ""
            pct = (
                int(int(current) / max(int(total), 1) * 100)
                if current.isdigit() and total.isdigit()
                else 0
            )
            label = f"{oq_label}: {pct}%"
            if eta:
                label += f" ({eta} remaining)"
            return label
        return labels.get(phase, phase)


# ===== FILE: omlx/admin/routes.py =====

# SPDX-License-Identifier: Apache-2.0
"""Admin panel routes for oMLX server configuration.

This module provides HTTP routes for the admin panel including:
- Login/logout with API key authentication
- Dashboard for server monitoring
- Model settings management (per-model sampling parameters, pinning, default)
- Global settings management
"""

import asyncio
import inspect
import json
import logging
import os
import re
import shutil
import signal
import sys
import time
from collections import deque
from dataclasses import asdict, is_dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal, Optional

import requests
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, Field

from ..api.markitdown import MARKITDOWN_MODEL_ID, markitdown_model_visible
from ..api.openai_models import _coerce_tool_call_arguments
from ..api.utils import _try_parse_json
from ..model_profiles import EXCLUDED_FROM_PROFILES
from ..model_settings import merge_chat_template_kwargs
from ..settings import BURST_DECODE_MODES, SubKeyEntry, burst_decode_env
from ..utils.release_check import normalize_update_channel, select_latest_release
from .auth import (
    REMEMBER_ME_MAX_AGE,
    SESSION_MAX_AGE,
    compare_keys,
    create_session_token,
    require_admin,
    validate_api_key,
    verify_api_key,
    verify_session,
)

logger = logging.getLogger(__name__)

PRESET_REMOTE_URL = "https://omlx.ai/assets/omlx_preset.json"


# =============================================================================
# Pydantic Models
# =============================================================================


class LoginRequest(BaseModel):
    """Request model for admin login."""

    api_key: str
    remember: bool = False


class SetupApiKeyRequest(BaseModel):
    """Request model for initial API key setup."""

    api_key: str
    api_key_confirm: str


class CreateSubKeyRequest(BaseModel):
    """Request model for creating a sub API key."""

    key: str
    name: str = ""


class DeleteSubKeyRequest(BaseModel):
    """Request model for deleting a sub API key."""

    key: str


class CacheProbeRequest(BaseModel):
    """Request model for probing per-prompt cache state.

    Tokenizes a chat message list with the target model's tokenizer, then
    classifies each block's location in the cache hierarchy:
    - Hot SSD (in-RAM copy of SSD cache, ready to mount without disk read)
    - Disk SSD (persisted only, needs disk read to reuse)
    - Cold (fully uncached — would require full prefill)
    """

    model_id: str
    messages: list[dict[str, Any]]
    tools: list[dict[str, Any]] | None = None
    chat_template_kwargs: dict[str, Any] | None = None
    thinking_budget: int | None = None


class ModelSettingsRequest(BaseModel):
    """Request model for updating per-model settings."""

    model_alias: str | None = None
    model_type_override: str | None = None
    max_context_window: int | None = None
    max_tokens: int | None = None
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    repetition_penalty: float | None = None
    min_p: float | None = None
    presence_penalty: float | None = None
    force_sampling: bool | None = None
    max_tool_result_tokens: int | None = None
    chat_template_kwargs: dict[str, Any] | None = None
    forced_ct_kwargs: list[str] | None = None
    ttl_seconds: int | None = None
    index_cache_freq: int | None = None
    enable_thinking: bool | None = None
    thinking_budget_enabled: bool | None = None
    thinking_budget_tokens: int | None = None
    # TurboQuant KV cache (mlx-vlm backend)
    turboquant_kv_enabled: bool | None = None
    turboquant_kv_bits: float | None = None
    # SpecPrefill (experimental)
    specprefill_enabled: bool | None = None
    specprefill_draft_model: str | None = None
    specprefill_keep_pct: float | None = None
    specprefill_threshold: int | None = None
    # DFlash (block diffusion speculative decoding)
    dflash_enabled: bool | None = None
    dflash_draft_model: str | None = None
    dflash_draft_quant_enabled: bool | None = None
    dflash_draft_quant_weight_bits: int | None = None
    dflash_draft_quant_activation_bits: int | None = None
    dflash_draft_quant_group_size: int | None = None
    dflash_max_ctx: int | None = None
    dflash_in_memory_cache: bool | None = None
    dflash_in_memory_cache_max_entries: int | None = None
    dflash_in_memory_cache_max_bytes: int | None = None
    dflash_ssd_cache: bool | None = None
    dflash_ssd_cache_max_bytes: int | None = None
    dflash_draft_window_size: int | None = None
    dflash_draft_sink_size: int | None = None
    dflash_verify_mode: str | None = None
    # Native MTP (mlx-lm PR 990 / PR 15 monkey-patch)
    mtp_enabled: bool | None = None
    # VLM MTP speculative decoding via external assistant drafter (mlx-vlm 191d7c8+)
    vlm_mtp_enabled: bool | None = None
    vlm_mtp_draft_model: str | None = None
    vlm_mtp_draft_block_size: int | None = None
    reasoning_parser: str | None = None
    guided_grammar_enabled: bool | None = None
    guided_grammar: str | None = None
    is_pinned: bool | None = None
    is_default: bool | None = None
    is_hidden: bool | None = None
    is_favorite: bool | None = None
    # Security: per-model opt-in for trust_remote_code (issue #926)
    trust_remote_code: bool | None = None


class CreateProfileRequest(BaseModel):
    """Request body for creating a per-model profile."""

    name: str
    display_name: str
    api_name: str | None = None
    description: str | None = None
    settings: dict[str, Any] = Field(default_factory=dict)
    also_save_as_template: bool = False
    source_template: str | None = None
    expose_as_model: bool = False


class UpdateProfileRequest(BaseModel):
    """Request body for updating/renaming a per-model profile."""

    new_name: str | None = None
    display_name: str | None = None
    api_name: str | None = None
    description: str | None = None
    settings: dict[str, Any] | None = None
    source_template: str | None = None
    expose_as_model: bool | None = None
    also_save_as_template: bool = False


class CreateTemplateRequest(BaseModel):
    """Request body for creating a global template."""

    name: str
    display_name: str
    description: str | None = None
    settings: dict[str, Any] = Field(default_factory=dict)


class UpdateTemplateRequest(BaseModel):
    """Request body for updating/renaming a global template."""

    new_name: str | None = None
    display_name: str | None = None
    description: str | None = None
    settings: dict[str, Any] | None = None


class GlobalSettingsRequest(BaseModel):
    """Request model for updating global server settings."""

    # Server settings
    host: str | None = None
    port: int | None = None
    log_level: str | None = None
    server_aliases: list[str] | None = None
    sse_keepalive_mode: str | None = None
    auto_start_on_launch: bool | None = None
    burst_decode_mode: str | None = None  # "off" / "light" / "balanced" / "aggressive"
    preserve_mid_system_cache: bool | None = None

    # Model settings
    model_dirs: list[str] | None = None
    model_dir: str | None = None  # Deprecated: kept for backward compatibility
    model_fallback: bool | None = None
    hide_helper_models: bool | None = None

    # Memory enforcement
    memory_prefill_memory_guard: bool | None = None
    memory_guard_tier: str | None = (
        None  # "safe" / "balanced" / "aggressive" / "custom"
    )
    memory_guard_custom_ceiling_gb: float | None = (
        None  # only used when tier == "custom"
    )

    # Scheduler settings
    max_concurrent_requests: int | None = None
    embedding_batch_size: int | None = None
    chunked_prefill: bool | None = None
    prefill_priority: str | None = None  # "context" | "speed"

    # Cache settings
    cache_enabled: bool | None = None
    ssd_cache_dir: str | None = None
    ssd_cache_max_size: str | None = None
    hot_cache_only: bool | None = None
    hot_cache_max_size: str | None = None  # "0" = disabled, "8GB", etc.
    initial_cache_blocks: int | None = None  # Starting blocks (requires restart)

    # MCP settings
    mcp_config: str | None = None

    # HuggingFace settings
    hf_endpoint: str | None = None
    hf_cache_enabled: bool | None = None

    # ModelScope settings
    ms_endpoint: str | None = None

    # Network settings
    network_http_proxy: str | None = None
    network_https_proxy: str | None = None
    network_no_proxy: str | None = None
    network_ca_bundle: str | None = None

    # Sampling defaults
    sampling_max_context_window: int | None = None
    sampling_max_context_window_policy: int | None = Field(default=None, ge=1)
    sampling_max_tokens: int | None = None
    sampling_temperature: float | None = None
    sampling_top_p: float | None = None
    sampling_top_k: int | None = None
    sampling_repetition_penalty: float | None = None

    # Claude Code settings
    claude_code_mode: str | None = None
    claude_code_opus_model: str | None = None
    claude_code_sonnet_model: str | None = None
    claude_code_haiku_model: str | None = None

    # Other integrations settings
    integrations_copilot_model: str | None = None
    integrations_codex_model: str | None = None
    integrations_opencode_model: str | None = None
    integrations_openclaw_model: str | None = None
    integrations_hermes_model: str | None = None
    integrations_pi_model: str | None = None
    integrations_openclaw_tools_profile: (
        Literal["minimal", "coding", "messaging", "full"] | None
    ) = None
    markitdown_enabled: bool | None = None
    markitdown_expose_model: bool | None = None
    markitdown_max_file_size_mb: int | None = None
    markitdown_max_files_per_request: int | None = None
    markitdown_pdf_processing_engine: str | None = None

    # UI settings
    ui_language: str | None = None

    # Idle timeout settings. null disables the global fallback.
    idle_timeout_seconds: int | None = Field(default=None, ge=60)

    # Auth settings
    api_key: str | None = None
    skip_api_key_verification: bool | None = None


class HFDownloadRequest(BaseModel):
    """Request model for starting a HuggingFace model download."""

    repo_id: str
    hf_token: str = ""


class HFRetryRequest(BaseModel):
    """Request model for retrying a HuggingFace model download."""

    hf_token: str = ""


class MSDownloadRequest(BaseModel):
    """Request model for starting a ModelScope model download."""

    model_id: str
    ms_token: str = ""


class MSRetryRequest(BaseModel):
    """Request model for retrying a ModelScope model download."""

    ms_token: str = ""


class OQStartRequest(BaseModel):
    """Request model for starting an oQ quantization task."""

    model_path: str
    oq_level: float
    group_size: int = 64
    sensitivity_model_path: str = ""
    text_only: bool = False
    dtype: str = "bfloat16"
    preserve_mtp: bool = False
    auto_proxy_sensitivity: bool = True
    enhanced: bool = False
    imatrix_cache_path: str = ""
    imatrix_reuse_cache: bool = True
    imatrix_strict: bool = False
    imatrix_num_samples: int = 128
    imatrix_seq_length: int = 512
    mtp_assistant_model_path: str = ""


class HFUploadRequest(BaseModel):
    """Request model for starting a HuggingFace upload task."""

    model_path: str
    repo_id: str
    hf_token: str
    readme_source_path: str = ""
    auto_readme: bool = True
    redownload_notice: bool = False
    private: bool = False


class HFValidateTokenRequest(BaseModel):
    """Request model for validating a HuggingFace token."""

    hf_token: str


# =============================================================================
# Runtime Settings Application Functions
# =============================================================================


def _format_cache_size(size_bytes: int) -> str:
    """Format cache size in bytes to human-readable string (e.g., '100GB')."""
    gb = size_bytes / (1024**3)
    if gb >= 1:
        return f"{gb:.0f}GB"
    mb = size_bytes / (1024**2)
    return f"{mb:.0f}MB"


def _parse_hot_cache_max_size(value: str) -> int:
    """Parse hot cache max size. Hot cache does not support an auto sentinel."""
    from ..config import parse_size

    normalized = value.strip()
    if normalized.lower() == "auto":
        raise ValueError(
            "Invalid hot_cache_max_size: 'auto' is not supported; "
            "use '0' to disable or a size like '8GB'"
        )

    try:
        size = parse_size(normalized)
    except ValueError as exc:
        raise ValueError(f"Invalid hot_cache_max_size: {exc}") from exc

    if size < 0:
        raise ValueError(
            "Invalid hot_cache_max_size: must be '0' to disable "
            "or a non-negative size"
        )
    return size


_PAROQUANT_REASON = "Not supported on paroquant models yet (compatibility not verified)"


def _paroquant_compat_for_model(model_info: dict) -> tuple[bool, str]:
    """Detect whether a model is paroquant-quantized.

    Returns ``(is_paroquant, reason)``. ``is_paroquant`` is True iff
    ``config.json`` declares ``quantization_config.quant_method == "paroquant"``.
    Reason is the user-facing string surfaced as a tooltip/banner on the
    admin model settings modal when paroquant gates an experimental toggle.
    """
    import json
    from pathlib import Path

    model_path = model_info.get("model_path") or ""
    if not model_path:
        return False, ""
    cfg_path = Path(model_path) / "config.json"
    if not cfg_path.exists():
        return False, ""
    try:
        cfg = json.loads(cfg_path.read_text())
    except Exception:
        return False, ""
    qcfg = cfg.get("quantization_config") or {}
    method = (qcfg.get("quant_method") or "").lower()
    if method == "paroquant":
        return True, _PAROQUANT_REASON
    return False, ""


def _dflash_compat_for_model(model_info: dict) -> tuple[bool, str]:
    """Resolve dflash compatibility for an engine_pool model dict.

    Returns ``(False, "")`` when dflash-mlx is not installed so the UI hides
    the compat hint instead of pointing the user at an unrelated reason.
    """
    is_paro, paro_reason = _paroquant_compat_for_model(model_info)
    if is_paro:
        return False, paro_reason
    try:
        from ..engine.dflash import is_dflash_compatible
    except ImportError:
        return False, ""
    model_path = model_info.get("model_path") or ""
    if not model_path:
        return False, "model_path missing"
    return is_dflash_compatible(model_path)


def _entry_is_diffusion_model(entry) -> bool:
    model_type = (getattr(entry, "config_model_type", None) or "").lower()
    return model_type.replace("-", "_") == "diffusion_gemma"


def _sanitize_diffusion_settings_dict(settings: dict) -> None:
    """Clear unsupported diffusion-lane settings before ModelSettings parsing.

    Tool-calling settings (``max_tool_result_tokens``) are intentionally NOT
    cleared: tool calling is prompt-driven plus output parsing and works on
    the diffusion lane when a tool parser matches the chat template.
    """
    unsupported_none_fields = (
        "top_p",
        "top_k",
        "min_p",
        "repetition_penalty",
        "presence_penalty",
        "enable_thinking",
        "preserve_thinking",
        "thinking_budget_tokens",
        "reasoning_parser",
        "guided_grammar",
        "index_cache_freq",
        "specprefill_draft_model",
        "specprefill_keep_pct",
        "specprefill_threshold",
        "dflash_draft_model",
        "dflash_draft_quant_enabled",
        "dflash_draft_quant_weight_bits",
        "dflash_draft_quant_activation_bits",
        "dflash_draft_quant_group_size",
        "dflash_max_ctx",
        "dflash_draft_window_size",
        "dflash_draft_sink_size",
        "dflash_verify_mode",
        "vlm_mtp_draft_model",
        "vlm_mtp_draft_block_size",
    )
    for key in unsupported_none_fields:
        settings[key] = None

    settings["force_sampling"] = False
    settings["thinking_budget_enabled"] = False
    settings["guided_grammar_enabled"] = False
    settings["turboquant_kv_enabled"] = False
    settings["turboquant_kv_bits"] = 4
    settings["turboquant_skip_last"] = True
    settings["specprefill_enabled"] = False
    settings["dflash_enabled"] = False
    settings["dflash_in_memory_cache"] = True
    settings["dflash_in_memory_cache_max_entries"] = 4
    settings["dflash_in_memory_cache_max_bytes"] = 8 * 1024 * 1024 * 1024
    settings["dflash_ssd_cache"] = False
    settings["dflash_ssd_cache_max_bytes"] = 20 * 1024 * 1024 * 1024
    settings["mtp_enabled"] = False
    settings["vlm_mtp_enabled"] = False

    unsupported_ct_kwargs = {
        "enable_thinking",
        "reasoning_effort",
        "preserve_thinking",
    }
    kwargs = settings.get("chat_template_kwargs")
    if kwargs:
        filtered_kwargs = {
            k: v for k, v in kwargs.items() if k not in unsupported_ct_kwargs
        }
        settings["chat_template_kwargs"] = filtered_kwargs or None
    forced = settings.get("forced_ct_kwargs")
    if forced:
        allowed = set(settings.get("chat_template_kwargs") or {})
        filtered_forced = [
            k for k in forced if k not in unsupported_ct_kwargs and k in allowed
        ]
        settings["forced_ct_kwargs"] = filtered_forced or None


def _sanitize_diffusion_model_settings(settings) -> None:
    """Clear settings that the serial diffusion lane does not implement.

    ``max_tool_result_tokens`` is intentionally preserved — tool calling
    works on the diffusion lane (prompt-driven + output parsing).
    """
    settings.top_p = None
    settings.top_k = None
    settings.min_p = None
    settings.repetition_penalty = None
    settings.presence_penalty = None
    settings.force_sampling = False
    settings.enable_thinking = None
    settings.preserve_thinking = None
    settings.thinking_budget_enabled = False
    settings.thinking_budget_tokens = None
    settings.reasoning_parser = None
    settings.guided_grammar_enabled = False
    settings.guided_grammar = None

    unsupported_ct_kwargs = {
        "enable_thinking",
        "reasoning_effort",
        "preserve_thinking",
    }
    if settings.chat_template_kwargs:
        filtered_kwargs = {
            k: v
            for k, v in settings.chat_template_kwargs.items()
            if k not in unsupported_ct_kwargs
        }
        settings.chat_template_kwargs = filtered_kwargs or None
    if settings.forced_ct_kwargs:
        allowed = set(settings.chat_template_kwargs or {})
        filtered_forced = [
            k
            for k in settings.forced_ct_kwargs
            if k not in unsupported_ct_kwargs and k in allowed
        ]
        settings.forced_ct_kwargs = filtered_forced or None

    settings.index_cache_freq = None
    settings.turboquant_kv_enabled = False
    settings.turboquant_kv_bits = 4
    settings.turboquant_skip_last = True
    settings.specprefill_enabled = False
    settings.specprefill_draft_model = None
    settings.specprefill_keep_pct = None
    settings.specprefill_threshold = None
    settings.dflash_enabled = False
    settings.dflash_draft_model = None
    settings.dflash_draft_quant_enabled = None
    settings.dflash_draft_quant_weight_bits = None
    settings.dflash_draft_quant_activation_bits = None
    settings.dflash_draft_quant_group_size = None
    settings.dflash_max_ctx = None
    settings.dflash_in_memory_cache = True
    settings.dflash_in_memory_cache_max_entries = 4
    settings.dflash_in_memory_cache_max_bytes = 8 * 1024 * 1024 * 1024
    settings.dflash_ssd_cache = False
    settings.dflash_ssd_cache_max_bytes = 20 * 1024 * 1024 * 1024
    settings.dflash_draft_window_size = None
    settings.dflash_draft_sink_size = None
    settings.dflash_verify_mode = None
    settings.mtp_enabled = False
    settings.vlm_mtp_enabled = False
    settings.vlm_mtp_draft_model = None
    settings.vlm_mtp_draft_block_size = None


def _mtp_compat_for_model(model_info: dict) -> tuple[bool, str]:
    """Mirror of ``_dflash_compat_for_model`` for the native MTP toggle.

    Returns ``(compatible, reason)``. Reason is empty on success and
    suitable for surfacing to users (admin UI shows it under the toggle).

    The check is conservative: even when the config declares MTP layers
    we also peek at the safetensors weight index to verify that the
    converter actually preserved the MTP tensors, using the loader's
    ``_checkpoint_has_mtp_weights`` so native nextn layouts
    (``model.layers.<num_hidden_layers + i>.*``, e.g. GLM-5.2) count as
    present (issue #2326). Default mlx-lm converters strip ``mtp.*``;
    PR 990 ships a separate path that keeps them.
    """
    import json
    from pathlib import Path

    from ..utils.model_loading import (
        _checkpoint_has_mtp_weights,
        _has_mtp_heads,
        _is_mtp_compatible,
    )

    is_paro, paro_reason = _paroquant_compat_for_model(model_info)
    if is_paro:
        return False, paro_reason

    model_path = model_info.get("model_path") or ""
    if not model_path:
        return False, "model_path missing"
    cfg_path = Path(model_path) / "config.json"
    if not cfg_path.exists():
        return False, "config.json not found"
    try:
        cfg = json.loads(cfg_path.read_text())
    except Exception as e:
        return False, f"failed to read config: {e}"
    model_type = cfg.get("model_type")
    if not _has_mtp_heads(cfg):
        return False, "model has no MTP heads in config"
    if not _is_mtp_compatible(cfg, model_type):
        return False, (
            f"model_type={model_type!r} is not on the MTP whitelist "
            "(supported: qwen3_5*, qwen3_6*, deepseek_v4*, glm_moe_dsa)"
        )
    if not _checkpoint_has_mtp_weights(model_path):
        return False, (
            "Config declares MTP layers but the weight files contain neither "
            "mtp.* tensors nor native nextn layers. Re-convert from HF with a "
            "converter that preserves MTP weights."
        )
    return True, ""


def _apply_log_level_runtime(level: str) -> None:
    """Apply log level change at runtime to all oMLX loggers and handlers."""
    level_name = level.upper()
    log_level = (
        5 if level_name == "TRACE" else getattr(logging, level_name, logging.INFO)
    )

    # Update root logger level and all its handlers
    root_logger = logging.getLogger()
    root_logger.setLevel(log_level)
    for handler in root_logger.handlers:
        handler.setLevel(log_level)

    # Update omlx-related loggers
    omlx_loggers = [
        "omlx",
        "omlx.scheduler",
        "omlx.paged_ssd_cache",
        "omlx.memory_monitor",
        "omlx.paged_cache",
        "omlx.prefix_cache",
        "omlx.engine_pool",
        "omlx.model_discovery",
        "omlx.engine_core",
        "omlx.engine",
        "omlx.server",
        "omlx.admin",
    ]

    for logger_name in omlx_loggers:
        logging.getLogger(logger_name).setLevel(log_level)

    # Also update uvicorn logger
    logging.getLogger("uvicorn").setLevel(log_level)
    logging.getLogger("uvicorn.access").setLevel(log_level)


async def _apply_model_dirs_runtime(model_dirs: list[str]) -> tuple[bool, str]:
    """
    Apply model directories change at runtime by re-scanning models.

    This will:
    1. Validate all directories
    2. Unload all currently loaded models
    3. Clear the entries dictionary
    4. Re-discover models from the new directories

    Returns:
        Tuple of (success, message)
    """
    from pathlib import Path

    from ..model_discovery import (
        model_directory_access_error,
        model_directory_write_error,
    )
    from ..server import _server_state

    if _server_state.engine_pool is None:
        return False, "Engine pool not initialized"

    if not model_dirs:
        return False, "At least one model directory is required"

    primary_path = Path(model_dirs[0]).expanduser().resolve()
    write_error = model_directory_write_error(primary_path, create=True)
    if write_error is not None:
        return False, write_error

    active_model_dirs = [str(primary_path)]
    for model_dir in model_dirs[1:]:
        model_path = Path(model_dir).expanduser().resolve()
        access_error = model_directory_access_error(model_path)
        if access_error is not None:
            logger.warning(
                "Skipping inaccessible model directory during runtime reload: %s",
                access_error,
            )
            continue
        active_model_dirs.append(str(model_path))

    pool = _server_state.engine_pool

    # Get pinned models from settings_manager
    pinned_models = []
    if _server_state.settings_manager is not None:
        pinned_models = _server_state.settings_manager.get_pinned_model_ids()

    # Unload all loaded models
    loaded_models = pool.get_loaded_model_ids()
    for model_id in loaded_models:
        try:
            await pool._unload_engine(model_id)
        except Exception as e:
            logger.warning(f"Error unloading {model_id}: {e}")

    # Clear entries
    pool._entries.clear()
    pool._current_model_memory = 0

    # Update downloader model directories
    global _hf_downloader, _ms_downloader, _oq_manager, _hf_uploader
    primary_dir = str(primary_path)
    if _hf_downloader is not None:
        _hf_downloader.update_model_dir(primary_dir)
    if _ms_downloader is not None:
        _ms_downloader.update_model_dir(primary_dir)

    # Update components that scan all model directories
    if _oq_manager is not None:
        _oq_manager.update_model_dirs(active_model_dirs)
    if _hf_uploader is not None:
        _hf_uploader.update_model_dirs(active_model_dirs)

    # Re-discover models from new directories
    try:
        pool.discover_models(active_model_dirs, pinned_models)
        if _server_state.settings_manager is not None:
            pool.apply_settings_overrides(_server_state.settings_manager)
    except Exception as e:
        return False, f"Failed to discover models: {e}"

    dir_count = len(active_model_dirs)
    return True, (
        f"Re-discovered {pool.model_count} models "
        f"from {dir_count} director{'ies' if dir_count > 1 else 'y'}"
    )


async def _reload_models() -> tuple[bool, str]:
    """
    Reload models: re-read model_settings.json, re-scan dirs, re-apply overrides,
    and preload pinned models.

    This does NOT re-read settings.json (global settings). It only refreshes
    the model inventory and per-model settings.

    Returns:
        Tuple of (success, message)
    """
    from ..server import _server_state

    if _server_state.engine_pool is None:
        return False, "Engine pool not initialized"

    global_settings = _get_global_settings()
    if global_settings is None:
        return False, "Global settings not initialized"

    # Re-read model_settings.json from disk
    settings_manager = _get_settings_manager()
    if settings_manager is not None:
        settings_manager._load()

    # Get current effective model dirs from global settings
    model_dirs = [str(d) for d in global_settings.get_effective_model_dirs()]

    # Unload all, re-discover, re-apply overrides
    success, msg = await _apply_model_dirs_runtime(model_dirs)
    if not success:
        return False, msg

    # Preload pinned models
    pool = _server_state.engine_pool
    if pool is not None:
        await pool.preload_pinned_models()

    return True, msg


async def _apply_memory_guard_tier_runtime(
    tier: str | None = None,
    custom_ceiling_gb: float | None = None,
) -> tuple[bool, str]:
    """
    Apply memory_guard_tier (and optionally custom ceiling) at runtime.

    Pushes both values into the running ProcessMemoryEnforcer, which
    recomputes static + dynamic ceilings on its next propagation tick.
    `tier` and `custom_ceiling_gb` can be passed together (Custom tier
    save) or independently.

    Returns:
        Tuple of (success, message)
    """
    from ..server import _server_state
    from ..settings import VALID_MEMORY_GUARD_TIERS

    enforcer = _server_state.process_memory_enforcer
    if enforcer is None:
        return False, "Process memory enforcer not initialized"

    changes = []
    if tier is not None:
        value = tier.strip().lower()
        if value not in VALID_MEMORY_GUARD_TIERS:
            return False, (
                f"Invalid memory_guard_tier: '{tier}' "
                f"(must be one of {sorted(VALID_MEMORY_GUARD_TIERS)})"
            )
        old_tier = enforcer.memory_guard_tier
        enforcer.memory_guard_tier = value
        changes.append(f"tier: {old_tier} -> {value}")
    if custom_ceiling_gb is not None:
        new_bytes = max(0, int(float(custom_ceiling_gb) * 1024**3))
        enforcer.memory_guard_custom_ceiling_bytes = new_bytes
        changes.append(f"custom_ceiling: {custom_ceiling_gb} GB")
    if not changes:
        return True, "(no change)"
    return True, "Memory guard updated — " + ", ".join(changes)


async def _apply_cache_settings_runtime(
    enabled: bool | None,
    ssd_cache_dir: str | None,
    ssd_cache_max_size: str | None,
    global_settings,
    hot_cache_max_size: str | None = None,
) -> tuple[bool, str]:
    """
    Apply cache settings at runtime.

    Updates the scheduler_config and unloads all models so they
    will use the new cache settings when reloaded.

    Returns:
        Tuple of (success, message)
    """
    from ..config import parse_size
    from ..server import _server_state

    if _server_state.engine_pool is None:
        return False, "Engine pool not initialized"

    pool = _server_state.engine_pool

    # Update scheduler config based on cache settings
    if enabled is False or (enabled is None and not global_settings.cache.enabled):
        pool._scheduler_config.paged_ssd_cache_dir = None
        pool._scheduler_config.paged_ssd_cache_max_size = 0
    else:
        # Cache is enabled
        if ssd_cache_dir is not None:
            pool._scheduler_config.paged_ssd_cache_dir = ssd_cache_dir
        elif global_settings.cache.ssd_cache_dir:
            pool._scheduler_config.paged_ssd_cache_dir = (
                global_settings.cache.ssd_cache_dir
            )
        else:
            # Use default cache dir
            pool._scheduler_config.paged_ssd_cache_dir = str(
                global_settings.cache.get_ssd_cache_dir(global_settings.base_path)
            )

        if ssd_cache_max_size is not None:
            # Handle "auto" value
            if ssd_cache_max_size.lower() == "auto":
                pool._scheduler_config.paged_ssd_cache_max_size = (
                    global_settings.cache.get_ssd_cache_max_size_bytes(
                        global_settings.base_path
                    )
                )
            else:
                pool._scheduler_config.paged_ssd_cache_max_size = parse_size(
                    ssd_cache_max_size
                )
        elif global_settings.cache.ssd_cache_max_size:
            # Use settings value (handles "auto")
            pool._scheduler_config.paged_ssd_cache_max_size = (
                global_settings.cache.get_ssd_cache_max_size_bytes(
                    global_settings.base_path
                )
            )
        elif global_settings.cache.ssd_cache_max_size:
            pool._scheduler_config.paged_ssd_cache_max_size = parse_size(
                global_settings.cache.ssd_cache_max_size
            )

    # Apply hot cache max size
    if hot_cache_max_size is not None:
        hot_bytes = _parse_hot_cache_max_size(hot_cache_max_size)
        old_hot = pool._scheduler_config.hot_cache_max_size
        pool._scheduler_config.hot_cache_max_size = hot_bytes
        if hot_bytes != old_hot:
            from ..utils.formatting import format_bytes

            old_str = "Off" if old_hot == 0 else format_bytes(old_hot)
            new_str = "Off" if hot_bytes == 0 else format_bytes(hot_bytes)
            logger.info(f"Hot cache max size changed: {old_str} -> {new_str}")
    elif global_settings.cache.hot_cache_max_size:
        pool._scheduler_config.hot_cache_max_size = (
            global_settings.cache.get_hot_cache_max_size_bytes()
        )
    if hasattr(pool, "configure_hot_cache_budget"):
        pool.configure_hot_cache_budget()

    # Unload all loaded models so they use new config when reloaded
    loaded_models = pool.get_loaded_model_ids()
    for model_id in loaded_models:
        try:
            await pool._unload_engine(model_id)
        except Exception as e:
            logger.warning(f"Error unloading {model_id}: {e}")

    return True, f"Cache settings updated. Unloaded {len(loaded_models)} models."


def _apply_sampling_settings_runtime(
    max_context_window: int | None,
    max_context_window_policy: int | None,
    max_context_window_policy_set: bool,
    max_tokens: int | None,
    temperature: float | None,
    top_p: float | None,
    top_k: int | None,
    repetition_penalty: float | None = None,
) -> tuple[bool, str]:
    """
    Apply sampling default settings at runtime.

    Updates _server_state.sampling which is used for all new API requests.

    Returns:
        Tuple of (success, message)
    """
    from ..server import _server_state

    changes = []

    if max_context_window is not None:
        _server_state.sampling.max_context_window = max_context_window
        changes.append(f"max_context_window={max_context_window}")

    if max_context_window_policy_set:
        _server_state.sampling.max_context_window_policy = max_context_window_policy
        changes.append(f"max_context_window_policy={max_context_window_policy}")

    if max_tokens is not None:
        _server_state.sampling.max_tokens = max_tokens
        changes.append(f"max_tokens={max_tokens}")

    if temperature is not None:
        _server_state.sampling.temperature = temperature
        changes.append(f"temperature={temperature}")

    if top_p is not None:
        _server_state.sampling.top_p = top_p
        changes.append(f"top_p={top_p}")

    if top_k is not None:
        _server_state.sampling.top_k = top_k
        changes.append(f"top_k={top_k}")

    if repetition_penalty is not None:
        _server_state.sampling.repetition_penalty = repetition_penalty
        changes.append(f"repetition_penalty={repetition_penalty}")

    if changes:
        return True, f"Sampling defaults updated: {', '.join(changes)}"
    return True, "No sampling changes"


# =============================================================================
# Router and Templates
# =============================================================================

router = APIRouter(prefix="/admin", tags=["admin"])
templates = Jinja2Templates(directory=Path(__file__).parent / "templates")
static_dir = Path(__file__).parent / "static"


def _static_version(path: str) -> str:
    """Append file mtime as query string for cache busting."""
    file_path = static_dir / path
    if file_path.is_file():
        mtime = int(file_path.stat().st_mtime)
        return f"/admin/static/{path}?v={mtime}"
    return f"/admin/static/{path}"


templates.env.globals["static"] = _static_version

from omlx._version import __version__ as _omlx_version

templates.env.globals["version"] = _omlx_version

# i18n defaults (English) — overridden once set_admin_getters is called
_i18n_dir = Path(__file__).parent / "i18n"
_en_locale: dict = {}
try:
    _en_locale = json.loads((_i18n_dir / "en.json").read_text(encoding="utf-8"))
except Exception:
    pass
templates.env.globals["t"] = lambda key: _en_locale.get(key, key)
templates.env.globals["locale_json"] = json.dumps(_en_locale, ensure_ascii=False)
templates.env.globals["current_lang"] = "en"


def _load_locale(language: str) -> dict:
    """Load locale dict and fill missing keys from English."""
    fallback = dict(_en_locale)
    path = _i18n_dir / f"{language}.json"
    if language == "en":
        return fallback
    try:
        locale = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        try:
            return json.loads((_i18n_dir / "en.json").read_text(encoding="utf-8"))
        except Exception:
            return {}
    fallback.update(locale)
    return fallback


def _make_t(locale: dict):
    """Return a Jinja2-compatible t() function for the given locale dict."""

    def t(key: str) -> str:
        return locale.get(key, key)

    return t


def _refresh_i18n_globals() -> None:
    """Reload i18n globals from current settings. Called on startup and language change."""
    lang = "en"
    try:
        settings = _get_global_settings() if _get_global_settings else None
        if settings:
            lang = settings.ui.language
    except Exception:
        pass
    locale = _load_locale(lang)
    templates.env.globals["t"] = _make_t(locale)
    templates.env.globals["locale_json"] = json.dumps(locale, ensure_ascii=False)
    templates.env.globals["current_lang"] = lang


# =============================================================================
# State Getters (set by server.py)
# =============================================================================

_get_server_state = None
_get_engine_pool = None
_get_settings_manager = None
_get_global_settings = None
_hf_downloader = None
_ms_downloader = None
_oq_manager = None
_hf_uploader = None


def set_admin_getters(
    state_getter,
    pool_getter,
    settings_manager_getter,
    global_settings_getter,
):
    """
    Set the getter functions for accessing server state.

    This function must be called during server initialization to provide
    access to the server state objects.

    Args:
        state_getter: Function that returns the ServerState instance.
        pool_getter: Function that returns the EnginePool instance.
        settings_manager_getter: Function that returns the ModelSettingsManager.
        global_settings_getter: Function that returns the GlobalSettings.
    """
    global _get_server_state, _get_engine_pool, _get_settings_manager, _get_global_settings
    _get_server_state = state_getter
    _get_engine_pool = pool_getter
    _get_settings_manager = settings_manager_getter
    _get_global_settings = global_settings_getter
    _refresh_i18n_globals()


def set_hf_downloader(downloader):
    """Set the HFDownloader instance for admin routes.

    Args:
        downloader: HFDownloader instance created during server initialization.
    """
    global _hf_downloader
    _hf_downloader = downloader


def set_ms_downloader(downloader):
    """Set the MSDownloader instance for admin routes.

    Args:
        downloader: MSDownloader instance created during server initialization.
    """
    global _ms_downloader
    _ms_downloader = downloader


def set_oq_manager(manager):
    """Set the OQManager instance for admin routes.

    Args:
        manager: OQManager instance created during server initialization.
    """
    global _oq_manager
    _oq_manager = manager


def set_hf_uploader(uploader):
    """Set the HFUploader instance for admin routes.

    Args:
        uploader: HFUploader instance created during server initialization.
    """
    global _hf_uploader
    _hf_uploader = uploader


# =============================================================================
# Helper Functions
# =============================================================================


def format_size(size_bytes: int) -> str:
    """
    Format a byte size as a human-readable string.

    Args:
        size_bytes: Size in bytes.

    Returns:
        Human-readable string (e.g., "1.5 GB").
    """
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024**2:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024**3:
        return f"{size_bytes / 1024**2:.1f} MB"
    elif size_bytes < 1024**4:
        return f"{size_bytes / 1024**3:.2f} GB"
    else:
        return f"{size_bytes / 1024**4:.2f} TB"


def get_ssd_disk_info(cache_dir: str) -> dict:
    """
    Get disk information for the SSD cache directory.

    Returns:
        Dictionary with total_bytes, total_formatted.
    """
    try:
        check_path = Path(cache_dir).expanduser().resolve()
        while not check_path.exists() and check_path.parent != check_path:
            check_path = check_path.parent
        stat = shutil.disk_usage(check_path)
        return {
            "total_bytes": stat.total,
            "total_formatted": format_size(stat.total),
        }
    except Exception as e:
        logger.warning(f"Failed to get disk info for {cache_dir}: {e}")
        return {
            "total_bytes": 0,
            "total_formatted": "Unknown",
        }


def get_system_memory_info() -> dict:
    """
    Get system memory information.

    Returns:
        Dictionary with total_bytes, total_formatted, auto_limit_bytes,
        and auto_limit_formatted (80% of total).
    """
    try:
        from ..utils import psutil_compat

        total_bytes = int(psutil_compat.get_total_memory())
    except Exception:
        total_bytes = 0

    auto_limit_bytes = int(total_bytes * 0.8)

    # Live values so the admin UI can preview the actual hard ceiling for any
    # tier (static_ceiling + dynamic_ceiling depend on these). Read on each
    # call — never cached.
    try:
        from ..utils import psutil_compat

        available_bytes = int(psutil_compat.virtual_memory().available)
    except Exception:
        available_bytes = 0
    try:
        from ..utils.proc_memory import get_phys_footprint

        omlx_phys_footprint_bytes = int(get_phys_footprint())
    except Exception:
        omlx_phys_footprint_bytes = 0

    # Effective Metal cap = sysctl iogpu.wired_limit_mb when set, else
    # Apple's max_recommended_working_set_size (~75% of RAM). The admin UI
    # compares this against the value oMLX wanted at start (static
    # ceiling) and warns when the cap is below the request.
    try:
        from ..process_memory_enforcer import get_effective_metal_cap_bytes

        iogpu_wired_limit_bytes = int(get_effective_metal_cap_bytes())
    except Exception:
        iogpu_wired_limit_bytes = 0
    omlx_wired_limit_request_bytes = 0
    try:
        from ..server import _server_state

        enforcer = getattr(_server_state, "process_memory_enforcer", None)
        if enforcer is not None:
            omlx_wired_limit_request_bytes = int(
                getattr(enforcer, "_metal_wired_limit_request", 0) or 0
            )
    except Exception:
        pass

    # Live macOS vm_stat layers so the admin dashboard can preview the
    # tier-aware ceiling (free + inactive + active * ratio). Zero on
    # non-macOS / call failure — JS falls back to available_bytes.
    free_memory_bytes = 0
    inactive_memory_bytes = 0
    active_memory_bytes = 0
    try:
        from ..utils import psutil_compat

        vm = psutil_compat.get_macos_vm_stats()
        if vm is not None:
            free_memory_bytes = int(vm.get("free", 0))
            inactive_memory_bytes = int(vm.get("inactive", 0))
            active_memory_bytes = int(vm.get("active", 0))
    except Exception:
        pass

    return {
        "total_bytes": total_bytes,
        "total_formatted": format_size(total_bytes),
        "auto_limit_bytes": auto_limit_bytes,
        "auto_limit_formatted": format_size(auto_limit_bytes),
        "available_bytes": available_bytes,
        "omlx_phys_footprint_bytes": omlx_phys_footprint_bytes,
        "iogpu_wired_limit_bytes": iogpu_wired_limit_bytes,
        "omlx_wired_limit_request_bytes": omlx_wired_limit_request_bytes,
        "free_memory_bytes": free_memory_bytes,
        "inactive_memory_bytes": inactive_memory_bytes,
        "active_memory_bytes": active_memory_bytes,
    }


# =============================================================================
# HTML Page Routes
# =============================================================================


@router.get("", response_class=HTMLResponse)
@router.get("/", response_class=HTMLResponse)
async def login_page(request: Request):
    """
    Render the admin login page or setup page.

    If no API key is configured, the page will show the initial setup form.
    Otherwise, it shows the standard login form.

    Returns:
        HTML login/setup page.
    """
    # Redirect to dashboard if already authenticated
    from .auth import verify_session

    if verify_session(request):
        return RedirectResponse(url="/admin/dashboard", status_code=302)

    global_settings = _get_global_settings()

    # Skip login page when skip_api_key_verification is enabled
    if global_settings is not None and global_settings.auth.skip_api_key_verification:
        return RedirectResponse(url="/admin/dashboard", status_code=302)

    api_key_configured = bool(global_settings and global_settings.auth.api_key)
    return templates.TemplateResponse(
        request,
        "login.html",
        {"api_key_configured": api_key_configured},
    )


@router.get("/dashboard", response_class=HTMLResponse)
async def dashboard_page(request: Request, is_admin: bool = Depends(require_admin)):
    """
    Render the admin dashboard page.

    Requires admin authentication via session cookie.

    Returns:
        HTML dashboard page with server status and model list.
    """
    return templates.TemplateResponse(request, "dashboard.html", {})


@router.get("/chat", response_class=HTMLResponse)
async def chat_page(request: Request, is_admin: bool = Depends(require_admin)):
    """
    Render the chat page for interacting with models.

    Requires admin authentication via session cookie.
    The API key is injected into the template context so that
    the chat page can auto-set it in localStorage, bypassing
    the manual API key entry modal.

    Returns:
        HTML chat page.
    """
    global_settings = _get_global_settings()
    api_key = global_settings.auth.api_key if global_settings else ""
    return templates.TemplateResponse(request, "chat.html", {"api_key": api_key or ""})


@router.get("/static/{path:path}")
async def admin_static(path: str):
    """Serve static files for admin panel (CSS, JS, fonts, logos, etc.)."""
    file_path = static_dir / path
    if not file_path.is_file() or not file_path.resolve().is_relative_to(
        static_dir.resolve()
    ):
        raise HTTPException(status_code=404, detail="File not found")
    media_types = {
        ".svg": "image/svg+xml",
        ".png": "image/png",
        ".ico": "image/x-icon",
        ".css": "text/css",
        ".js": "application/javascript",
        ".woff2": "font/woff2",
        ".woff": "font/woff",
        ".ttf": "font/ttf",
    }
    media_type = media_types.get(file_path.suffix, "application/octet-stream")
    return FileResponse(file_path, media_type=media_type)


# =============================================================================
# Authentication API Routes
# =============================================================================


@router.post("/api/login")
async def login(request: LoginRequest, response: Response):
    """
    Authenticate with API key and create session.

    Requires an API key to be configured on the server. If no API key
    is configured, returns 400 directing the user to set one up first.

    Args:
        request: LoginRequest containing the API key.
        response: FastAPI response object for setting cookies.

    Returns:
        JSON response with success status.

    Raises:
        HTTPException: 400 if no API key configured, 401 if invalid.
    """
    global_settings = _get_global_settings()
    server_api_key = global_settings.auth.api_key if global_settings else None

    # Reject login if no API key is configured (must use setup first)
    if not server_api_key:
        raise HTTPException(
            status_code=400,
            detail="No API key configured. Please set up an API key first.",
        )

    # Main key only — sub keys must not grant admin login
    if not verify_api_key(request.api_key, server_api_key):
        raise HTTPException(
            status_code=401,
            detail="Invalid API key",
        )

    # Create session token and set cookie
    token = create_session_token(remember=request.remember)
    cookie_max_age = REMEMBER_ME_MAX_AGE if request.remember else SESSION_MAX_AGE
    response.set_cookie(
        key="omlx_admin_session",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=cookie_max_age,
    )

    return {"success": True}


@router.post("/api/setup-api-key")
async def setup_api_key(request: SetupApiKeyRequest, response: Response):
    """
    Set up the initial API key when none is configured.

    This endpoint is only available when no API key is currently set.
    After successful setup, a session is created so the user is
    immediately logged in.

    Args:
        request: SetupApiKeyRequest with api_key and api_key_confirm.
        response: FastAPI response object for setting cookies.

    Returns:
        JSON response with success status.

    Raises:
        HTTPException: 400 if key already configured, validation fails,
                      or keys don't match.
    """
    from ..server import _server_state

    global_settings = _get_global_settings()

    # Only allow setup if no API key is currently configured
    if global_settings and global_settings.auth.api_key:
        raise HTTPException(
            status_code=400,
            detail="API key is already configured. Use settings to change it.",
        )

    # Validate confirmation match
    if request.api_key != request.api_key_confirm:
        raise HTTPException(status_code=400, detail="API keys do not match")

    # Validate key format
    is_valid, error_msg = validate_api_key(request.api_key)
    if not is_valid:
        raise HTTPException(status_code=400, detail=error_msg)

    # Apply to settings and runtime
    global_settings.auth.api_key = request.api_key
    _server_state.api_key = request.api_key

    # Persist to file
    try:
        global_settings.save()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to save settings: {e}")

    logger.info("API key configured via initial setup")

    # Create session token and set cookie (auto-login after setup)
    token = create_session_token()
    response.set_cookie(
        key="omlx_admin_session",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=86400,  # 24 hours
    )

    return {"success": True, "message": "API key configured successfully"}


@router.post("/api/logout")
async def logout(response: Response):
    """
    Clear session cookie and logout.

    Args:
        response: FastAPI response object for clearing cookies.

    Returns:
        JSON response with success status.
    """
    response.delete_cookie(key="omlx_admin_session")
    return {"success": True}


@router.get("/auto-login")
async def auto_login(key: str = "", redirect: str = "/admin/dashboard"):
    """
    Auto-login using API key and redirect to the target admin page.

    Used by the macOS menubar app to open admin pages with automatic
    authentication, bypassing the manual login form.

    Args:
        key: The API key for authentication.
        redirect: The path to redirect to after login. Must start with /admin.

    Returns:
        HTTP 302 redirect with session cookie set.
    """
    if not redirect.startswith("/admin"):
        raise HTTPException(status_code=400, detail="Invalid redirect path")

    global_settings = _get_global_settings()
    server_api_key = global_settings.auth.api_key if global_settings else None

    # Main key only — sub keys must not grant admin login
    if not key or not server_api_key or not verify_api_key(key, server_api_key):
        return RedirectResponse(url="/admin", status_code=302)

    token = create_session_token()
    response = RedirectResponse(url=redirect, status_code=302)
    response.set_cookie(
        key="omlx_admin_session",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=86400,
    )
    return response


# =============================================================================
# Sub Key Management Routes
# =============================================================================


@router.post("/api/sub-keys")
async def create_sub_key(
    request: CreateSubKeyRequest, is_admin: bool = Depends(require_admin)
):
    """Create a new sub API key.

    Sub keys can only be used for API authentication, not admin login.

    Args:
        request: CreateSubKeyRequest with key and optional name.

    Returns:
        JSON with the created sub key entry.

    Raises:
        HTTPException: 400 if validation fails or key already exists.
    """
    global_settings = _get_global_settings()
    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Validate key format
    is_valid, error_msg = validate_api_key(request.key)
    if not is_valid:
        raise HTTPException(status_code=400, detail=error_msg)

    # Check for duplicate (against main key and existing sub keys)
    if global_settings.auth.api_key and compare_keys(
        request.key, global_settings.auth.api_key
    ):
        raise HTTPException(
            status_code=400, detail="Sub key cannot be the same as the main key"
        )

    for sk in global_settings.auth.sub_keys:
        if sk.key and compare_keys(request.key, sk.key):
            raise HTTPException(status_code=400, detail="This key already exists")

    entry = SubKeyEntry(
        key=request.key,
        name=request.name or "",
        created_at=datetime.now(UTC).isoformat(),
    )
    global_settings.auth.sub_keys.append(entry)

    try:
        global_settings.save()
    except Exception as e:
        # Rollback
        global_settings.auth.sub_keys.pop()
        raise HTTPException(status_code=500, detail=f"Failed to save settings: {e}")

    logger.info(f"Sub key created: {request.name or '(unnamed)'}")
    return {"success": True, "sub_key": entry.to_dict()}


@router.delete("/api/sub-keys")
async def delete_sub_key(
    request: DeleteSubKeyRequest, is_admin: bool = Depends(require_admin)
):
    """Delete a sub API key.

    Args:
        request: DeleteSubKeyRequest with the key to delete.

    Returns:
        JSON with success status.

    Raises:
        HTTPException: 404 if key not found.
    """
    global_settings = _get_global_settings()
    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Find and remove the key
    for i, sk in enumerate(global_settings.auth.sub_keys):
        if sk.key and compare_keys(request.key, sk.key):
            removed = global_settings.auth.sub_keys.pop(i)
            try:
                global_settings.save()
            except Exception as e:
                global_settings.auth.sub_keys.insert(i, removed)
                raise HTTPException(
                    status_code=500, detail=f"Failed to save settings: {e}"
                )
            logger.info(f"Sub key deleted: {sk.name or '(unnamed)'}")
            return {"success": True}

    raise HTTPException(status_code=404, detail="Sub key not found")


# =============================================================================
# Grammar API Routes
# =============================================================================


_SUPPORTED_MODELS_DOC_RE = re.compile(
    r"Supported models:\s*\n((?:\s*-\s*\S.*\n?)+)",
)


def _models_from_docstring(fn) -> list[str]:
    """Extract the ``Supported models:`` bullet list from an xgrammar 0.1.34+
    structural-tag function's docstring. Returns ``[]`` if the section is
    absent or unparseable."""
    doc = inspect.getdoc(fn) or ""
    match = _SUPPORTED_MODELS_DOC_RE.search(doc)
    if not match:
        return []
    return [
        line.strip().lstrip("-").strip()
        for line in match.group(1).splitlines()
        if line.strip().startswith("-")
    ]


@router.get("/api/grammar/parsers")
async def list_grammar_parsers(is_admin: bool = Depends(require_admin)):
    """Return available reasoning parser names from xgrammar.

    Supports both API generations:

    - **xgrammar 0.1.34+** exposes a per-model registry at
      ``xgrammar.builtin_structural_tag._structural_tag_registry``; supported
      model names are pulled from each function's docstring.
    - **xgrammar 0.1.32–0.1.33** exposes the now-removed helper
      ``get_builtin_structural_tag_supported_models()``.

    Returns ``[]`` if xgrammar is missing, fails to load (e.g. broken native
    binding on macOS arm64), or has neither API available.
    """
    # Install the torch stub BEFORE any xgrammar import. If this lives
    # inside the first try-block, a failure on the 0.1.34+ path can leave
    # the fallback try-block importing xgrammar without the stub, which
    # is guaranteed ImportError on stub-only (DMG) deployments.
    try:
        from omlx._torch_stub import install as _install_torch_stub

        _install_torch_stub()
    except Exception as e:  # pragma: no cover — defensive
        logger.debug("torch stub install failed: %s", e)

    # Prefer the 0.1.34+ registry so newer parsers (qwen3_6, gemma4,
    # deepseek_v4, ...) are exposed.
    try:
        from xgrammar.builtin_structural_tag import _structural_tag_registry

        return [
            {"value": style, "label": style, "models": _models_from_docstring(fn)}
            for style, fn in _structural_tag_registry.items()
        ]
    except Exception as e:
        logger.debug("xgrammar 0.1.34+ registry unavailable: %s", e)

    # Fall back to the pre-0.1.34 helper.
    try:
        from xgrammar import get_builtin_structural_tag_supported_models

        supported = get_builtin_structural_tag_supported_models()
        return [
            {"value": style, "label": style, "models": models}
            for style, models in supported.items()
        ]
    except Exception as e:
        logger.warning("xgrammar parser discovery unavailable: %s", e)
        return []


# =============================================================================
# Models API Routes
# =============================================================================


def _model_display_name(
    model_id: str,
    model_path: str | Path | None,
    model_dirs: list[Path],
    *,
    source_repo_id: str | None = None,
) -> str:
    """Return the UI-only display name for a discovered local model."""
    repo_id = (source_repo_id or "").strip()
    if "/" in repo_id:
        return repo_id

    if not model_path:
        return model_id

    path_text = str(model_path)
    if "://" in path_text:
        return model_id

    try:
        path = Path(path_text).expanduser().resolve()
    except (OSError, RuntimeError):
        path = Path(path_text).expanduser()

    for model_dir in model_dirs:
        try:
            rel = path.relative_to(model_dir.expanduser().resolve())
        except (OSError, RuntimeError, ValueError):
            continue

        parts = rel.parts
        if len(parts) >= 2:
            return f"{parts[0]}/{parts[1]}"
        return model_id

    return model_id


def _model_dirs_for_display(global_settings: Any | None) -> list[Path]:
    if global_settings is None:
        return []
    try:
        return global_settings.model.get_model_dirs(global_settings.base_path)
    except Exception as e:  # pragma: no cover - defensive for partial test doubles
        logger.debug("Could not resolve model dirs for display names: %s", e)
        return []


@router.get("/api/models")
async def list_models(is_admin: bool = Depends(require_admin)):
    """
    List all models with their settings.

    Returns model information from the engine pool combined with
    per-model settings from the settings manager.

    Returns:
        JSON list of models with their status and settings.

    Raises:
        HTTPException: 401 if not authenticated, 503 if server not initialized.
    """
    engine_pool = _get_engine_pool()
    settings_manager = _get_settings_manager()
    server_state = _get_server_state()
    global_settings = _get_global_settings() if _get_global_settings else None
    model_dirs = _model_dirs_for_display(global_settings)

    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Get engine pool status
    status = engine_pool.get_status()
    models_status = status.get("models", [])

    # Get all model settings
    all_settings = settings_manager.get_all_settings() if settings_manager else {}

    # Draft-model references pointed at by other models' speculative settings —
    # used to badge "helper" drafters that only differ by being referenced.
    referenced_drafts: set[str] = set()
    for _ms in all_settings.values():
        for ref in (
            _ms.specprefill_draft_model,
            _ms.dflash_draft_model,
            _ms.vlm_mtp_draft_model,
        ):
            if ref:
                referenced_drafts.add(ref)

    # SSD cache dir is set on the scheduler_config when the user enables paged
    # SSD caching; admin UI consumes it to gate the dflash SSD toggle.
    ssd_cache_dir = getattr(
        getattr(engine_pool, "_scheduler_config", None),
        "paged_ssd_cache_dir",
        None,
    )
    dflash_ssd_cache_available = bool(ssd_cache_dir)

    # Combine model info with settings
    models = []
    for model_info in models_status:
        model_id = model_info["id"]
        settings = all_settings.get(model_id)

        is_paroquant, paroquant_reason = _paroquant_compat_for_model(model_info)
        compat_ok, compat_reason = _dflash_compat_for_model(model_info)
        mtp_compat_ok, mtp_compat_reason = _mtp_compat_for_model(model_info)

        model_data = {
            "id": model_id,
            "model_path": model_info.get("model_path", ""),
            "display_name": _model_display_name(
                model_id,
                model_info.get("model_path", ""),
                model_dirs,
                source_repo_id=model_info.get("source_repo_id"),
            ),
            "loaded": model_info.get("loaded", False),
            "is_loading": model_info.get("is_loading", False),
            "estimated_size": model_info.get("estimated_size", 0),
            "estimated_size_formatted": format_size(
                model_info.get("estimated_size", 0)
            ),
            "actual_size": model_info.get("actual_size") or 0,
            "actual_size_formatted": (
                format_size(model_info.get("actual_size", 0))
                if model_info.get("actual_size")
                else None
            ),
            "pinned": model_info.get("pinned", False),
            "is_default": (
                server_state.default_model == model_id if server_state else False
            ),
            "is_hidden": bool(settings and settings.is_hidden),
            "is_favorite": bool(settings and settings.is_favorite),
            "is_helper": (
                bool(model_info.get("is_helper"))
                or model_id in referenced_drafts
                or model_info.get("model_path") in referenced_drafts
                or model_info.get("source_repo_id") in referenced_drafts
            ),
            "engine_type": model_info.get("engine_type", "batched"),
            "model_type": model_info.get("model_type", "llm"),
            "config_model_type": model_info.get("config_model_type", ""),
            # Native context window from the model's config.json — used by
            # the context bench UI to hide targets the model cannot reach.
            "model_context_length": model_info.get("model_context_length"),
            "thinking_default": model_info.get("thinking_default"),
            "preserve_thinking_default": model_info.get("preserve_thinking_default"),
            "source_type": model_info.get("source_type", "local"),
            "source_repo_id": model_info.get("source_repo_id"),
            "last_access": model_info.get("last_access"),
            "dflash_compatible": compat_ok,
            "dflash_compatibility_reason": compat_reason,
            "dflash_ssd_cache_available": dflash_ssd_cache_available,
            "mtp_compatible": mtp_compat_ok,
            "mtp_compatibility_reason": mtp_compat_reason,
            "is_paroquant": is_paroquant,
            "paroquant_reason": paroquant_reason,
        }

        # Add settings if available
        if settings:
            model_data["settings"] = asdict(settings)
        if settings_manager:
            model_data["exposed_profiles"] = [
                profile
                for profile in settings_manager.list_profiles(model_id)
                if profile.get("expose_as_model")
            ]

        models.append(model_data)

    if markitdown_model_visible(global_settings) and not any(
        m.get("id") == MARKITDOWN_MODEL_ID for m in models
    ):
        models.append(
            {
                "id": MARKITDOWN_MODEL_ID,
                "model_path": "builtin://markitdown",
                "display_name": MARKITDOWN_MODEL_ID,
                "loaded": True,
                "is_loading": False,
                "estimated_size": 0,
                "estimated_size_formatted": format_size(0),
                "actual_size": 0,
                "actual_size_formatted": None,
                "pinned": False,
                "is_default": False,
                "engine_type": "markitdown",
                "model_type": "markitdown",
                "config_model_type": "markitdown",
                "thinking_default": None,
                "preserve_thinking_default": None,
                "source_type": "builtin",
                "source_repo_id": None,
                "last_access": None,
                "dflash_compatible": False,
                "dflash_compatibility_reason": "",
                "dflash_ssd_cache_available": False,
                "mtp_compatible": False,
                "mtp_compatibility_reason": "",
                "is_paroquant": False,
                "paroquant_reason": "",
                "virtual": True,
            }
        )

    return {"models": models}


@router.post("/api/models/{model_id}/unload")
async def unload_model(
    model_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Manually unload a model from memory."""
    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    entry = engine_pool.get_entry(model_id)
    if entry is None:
        raise HTTPException(status_code=404, detail=f"Model not found: {model_id}")
    if entry.engine is None:
        raise HTTPException(status_code=400, detail=f"Model not loaded: {model_id}")

    await engine_pool._unload_engine(model_id)
    logger.info(f"Manually unloaded model: {model_id}")
    return {"status": "ok", "model_id": model_id, "message": f"Unloaded {model_id}"}


async def _require_admin_or_bearer(request: Request) -> bool:
    """Allow admin session OR a valid Bearer API key (for CLI use)."""
    gs = _get_global_settings() if _get_global_settings else None

    # No-auth mode: always allow
    if gs is not None and gs.auth.skip_api_key_verification:
        return True

    # Valid admin session cookie
    if verify_session(request):
        return True

    # Bearer token matching the configured API key
    auth_header = request.headers.get("Authorization", "")
    if auth_header.startswith("Bearer ") and gs is not None:
        token = auth_header[7:]
        server_key = gs.auth.api_key or ""
        sub_keys = gs.auth.sub_keys or []
        if verify_api_key(token, server_key):
            return True
        for sk in sub_keys:
            if verify_api_key(token, getattr(sk, "key", "")):
                return True

    raise HTTPException(
        status_code=401,
        detail="Admin authentication required",
        headers={"WWW-Authenticate": "Bearer"},
    )


@router.post("/api/models/{model_id}/load")
async def load_model(
    model_id: str,
    is_admin: bool = Depends(_require_admin_or_bearer),
):
    """Manually load a model into memory."""
    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    entry = engine_pool.get_entry(model_id)
    if entry is None:
        raise HTTPException(status_code=404, detail=f"Model not found: {model_id}")
    if entry.engine is not None:
        return {
            "status": "ok",
            "model_id": model_id,
            "message": f"Already loaded: {model_id}",
        }
    if entry.is_loading:
        raise HTTPException(
            status_code=409, detail=f"Model is already loading: {model_id}"
        )

    try:
        await engine_pool.get_engine(model_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    logger.info(f"Manually loaded model: {model_id}")
    return {"status": "ok", "model_id": model_id, "message": f"Loaded {model_id}"}


@router.post("/api/reload")
async def reload_models(is_admin: bool = Depends(require_admin)):
    """Reload models: re-read model settings, re-discover models, preload pinned."""
    success, message = await _reload_models()
    if success:
        return {"status": "ok", "message": message}
    raise HTTPException(status_code=500, detail=message)


@router.put("/api/models/{model_id}/settings")
async def update_model_settings(
    model_id: str,
    request: ModelSettingsRequest,
    is_admin: bool = Depends(require_admin),
):
    """
    Update settings for a specific model.

    Updates are persisted to the settings file and applied immediately
    to the engine pool where applicable (e.g., pinned status).

    Args:
        model_id: The model identifier.
        request: ModelSettingsRequest with the new settings.

    Returns:
        JSON response with success status and updated settings.

    Raises:
        HTTPException: 401 if not authenticated, 404 if model not found.
    """
    engine_pool = _get_engine_pool()
    settings_manager = _get_settings_manager()
    server_state = _get_server_state()

    if engine_pool is None or settings_manager is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Check if model exists
    entry = engine_pool.get_entry(model_id)
    if entry is None:
        raise HTTPException(status_code=404, detail=f"Model not found: {model_id}")

    # Get current settings
    current_settings = settings_manager.get_settings(model_id)

    # Apply updates — use model_fields_set to distinguish "sent as null"
    # (clear to default) from "not sent" (don't touch).
    sent = request.model_fields_set
    prev_engine_type = entry.engine_type  # Track for requires_reload check
    is_diffusion_model = _entry_is_diffusion_model(entry)
    if "model_alias" in sent:
        alias_value = request.model_alias.strip() if request.model_alias else None
        if alias_value == "":
            alias_value = None
        if alias_value is not None:
            all_settings = settings_manager.get_all_settings()
            for mid, ms in all_settings.items():
                if mid != model_id and ms.model_alias == alias_value:
                    raise HTTPException(
                        status_code=400,
                        detail=f"Alias '{alias_value}' is already used by model '{mid}'",
                    )
            for mid in engine_pool._entries:
                if mid != model_id and mid == alias_value:
                    raise HTTPException(
                        status_code=400,
                        detail=f"Alias '{alias_value}' conflicts with model directory name '{mid}'",
                    )
            _raise_if_alias_conflicts_exposed_profiles(
                alias_value=alias_value,
                model_id=model_id,
                settings_manager=settings_manager,
                engine_pool=engine_pool,
            )
        current_settings.model_alias = alias_value
    if "model_type_override" in sent:
        valid_types = {
            "llm",
            "vlm",
            "embedding",
            "reranker",
            "audio_stt",
            "audio_tts",
            "audio_sts",
        }
        # Treat empty string as None (auto-detect)
        override_value = request.model_type_override or None
        if override_value is not None and override_value not in valid_types:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid model_type_override: {request.model_type_override}",
            )
        current_settings.model_type_override = override_value
        # Update engine pool entry type immediately
        type_to_engine = {
            "llm": "batched",
            "vlm": "vlm",
            "embedding": "embedding",
            "reranker": "reranker",
            "audio_stt": "audio_stt",
            "audio_tts": "audio_tts",
            "audio_sts": "audio_sts",
        }
        if override_value:
            entry.model_type = override_value
            entry.engine_type = type_to_engine.get(override_value, "batched")
        else:
            # Reset to auto-detected type
            from pathlib import Path

            from ..model_discovery import detect_model_type

            detected_type = detect_model_type(Path(entry.model_path))
            entry.model_type = detected_type
            entry.engine_type = type_to_engine.get(detected_type, "batched")
    if "max_context_window" in sent:
        current_settings.max_context_window = request.max_context_window
    if "max_tokens" in sent:
        current_settings.max_tokens = request.max_tokens
    if "temperature" in sent:
        current_settings.temperature = request.temperature
    if "top_p" in sent:
        current_settings.top_p = request.top_p
    if "top_k" in sent:
        current_settings.top_k = request.top_k
    if "repetition_penalty" in sent:
        current_settings.repetition_penalty = request.repetition_penalty
    if "min_p" in sent:
        current_settings.min_p = request.min_p
    if "presence_penalty" in sent:
        current_settings.presence_penalty = request.presence_penalty
    if "force_sampling" in sent:
        current_settings.force_sampling = request.force_sampling
    if "max_tool_result_tokens" in sent:
        # 0 means disable (reset to None)
        current_settings.max_tool_result_tokens = (
            request.max_tool_result_tokens
            if request.max_tool_result_tokens and request.max_tool_result_tokens > 0
            else None
        )
    if "enable_thinking" in sent:
        current_settings.enable_thinking = request.enable_thinking
    if "thinking_budget_enabled" in sent:
        current_settings.thinking_budget_enabled = (
            request.thinking_budget_enabled or False
        )
    if "thinking_budget_tokens" in sent:
        current_settings.thinking_budget_tokens = (
            request.thinking_budget_tokens
            if request.thinking_budget_tokens and request.thinking_budget_tokens > 0
            else None
        )
    if "chat_template_kwargs" in sent:
        current_settings.chat_template_kwargs = request.chat_template_kwargs
    if "forced_ct_kwargs" in sent:
        current_settings.forced_ct_kwargs = request.forced_ct_kwargs
    if "ttl_seconds" in sent:
        current_settings.ttl_seconds = request.ttl_seconds
    if "index_cache_freq" in sent:
        # 0 means disable (reset to None)
        current_settings.index_cache_freq = (
            request.index_cache_freq
            if request.index_cache_freq and request.index_cache_freq >= 2
            else None
        )
    # TurboQuant KV cache settings
    if "turboquant_kv_enabled" in sent:
        current_settings.turboquant_kv_enabled = request.turboquant_kv_enabled or False
    if "turboquant_kv_bits" in sent:
        current_settings.turboquant_kv_bits = request.turboquant_kv_bits or 4
    # SpecPrefill settings
    if "specprefill_enabled" in sent:
        current_settings.specprefill_enabled = request.specprefill_enabled or False
    if "specprefill_draft_model" in sent:
        current_settings.specprefill_draft_model = (
            request.specprefill_draft_model or None
        )
    if "specprefill_keep_pct" in sent:
        current_settings.specprefill_keep_pct = request.specprefill_keep_pct or None
    if "specprefill_threshold" in sent:
        current_settings.specprefill_threshold = request.specprefill_threshold or None
    # DFlash settings
    if "dflash_enabled" in sent:
        new_dflash_enabled = (
            False if is_diffusion_model else bool(request.dflash_enabled)
        )
        if new_dflash_enabled:
            from ..engine.dflash import is_dflash_compatible

            compat_ok, compat_reason = is_dflash_compatible(entry.model_path)
            if not compat_ok:
                raise HTTPException(status_code=400, detail=compat_reason)
        current_settings.dflash_enabled = new_dflash_enabled
    if "dflash_draft_model" in sent:
        current_settings.dflash_draft_model = request.dflash_draft_model or None
    if "dflash_draft_quant_enabled" in sent:
        current_settings.dflash_draft_quant_enabled = (
            bool(request.dflash_draft_quant_enabled)
            if request.dflash_draft_quant_enabled is not None
            else None
        )
    if "dflash_draft_quant_weight_bits" in sent:
        current_settings.dflash_draft_quant_weight_bits = (
            int(request.dflash_draft_quant_weight_bits)
            if request.dflash_draft_quant_weight_bits is not None
            else None
        )
    if "dflash_draft_quant_activation_bits" in sent:
        current_settings.dflash_draft_quant_activation_bits = (
            int(request.dflash_draft_quant_activation_bits)
            if request.dflash_draft_quant_activation_bits is not None
            else None
        )
    if "dflash_draft_quant_group_size" in sent:
        current_settings.dflash_draft_quant_group_size = (
            int(request.dflash_draft_quant_group_size)
            if request.dflash_draft_quant_group_size is not None
            else None
        )
    if "dflash_max_ctx" in sent:
        # 0/None means "unlimited" — the engine treats None as no fallback threshold
        value = request.dflash_max_ctx
        current_settings.dflash_max_ctx = value if value and value > 0 else None
    if "dflash_in_memory_cache" in sent:
        current_settings.dflash_in_memory_cache = bool(request.dflash_in_memory_cache)
    if "dflash_in_memory_cache_max_entries" in sent:
        value = request.dflash_in_memory_cache_max_entries
        current_settings.dflash_in_memory_cache_max_entries = (
            int(value) if value and value > 0 else 4
        )
    if (
        "dflash_in_memory_cache_max_bytes" in sent
        and request.dflash_in_memory_cache_max_bytes
    ):
        current_settings.dflash_in_memory_cache_max_bytes = int(
            request.dflash_in_memory_cache_max_bytes
        )
    if "dflash_ssd_cache" in sent:
        ssd_requested = bool(request.dflash_ssd_cache)
        if is_diffusion_model:
            ssd_requested = False
        elif ssd_requested:
            in_mem_after = (
                bool(request.dflash_in_memory_cache)
                if "dflash_in_memory_cache" in sent
                else current_settings.dflash_in_memory_cache
            )
            if not in_mem_after:
                raise HTTPException(
                    status_code=400,
                    detail="DFlash SSD cache requires the in-memory cache to be enabled.",
                )
            ssd_dir = getattr(
                getattr(_get_engine_pool(), "_scheduler_config", None),
                "paged_ssd_cache_dir",
                None,
            )
            if not ssd_dir:
                raise HTTPException(
                    status_code=400,
                    detail=(
                        "DFlash SSD cache requires oMLX paged SSD cache to be enabled "
                        "(set --paged-ssd-cache-dir or configure it in settings)."
                    ),
                )
        current_settings.dflash_ssd_cache = ssd_requested
    if "dflash_ssd_cache_max_bytes" in sent and request.dflash_ssd_cache_max_bytes:
        current_settings.dflash_ssd_cache_max_bytes = int(
            request.dflash_ssd_cache_max_bytes
        )
    if "dflash_draft_window_size" in sent:
        # 0 / None / negative → fall back to dflash-mlx internal default (1024).
        value = request.dflash_draft_window_size
        current_settings.dflash_draft_window_size = (
            int(value) if value and value > 0 else None
        )
    if "dflash_draft_sink_size" in sent:
        # Negative is invalid; 0 is a legal sink-size (no sink tokens).
        value = request.dflash_draft_sink_size
        current_settings.dflash_draft_sink_size = (
            int(value) if value is not None and value >= 0 else None
        )
    if "dflash_verify_mode" in sent:
        value = request.dflash_verify_mode
        # dflash-mlx accepts: dflash | adaptive | ddtree | off.
        # Anything else (including empty string) → revert to dflash default.
        current_settings.dflash_verify_mode = (
            value if value in ("dflash", "adaptive", "ddtree", "off") else None
        )

    # Native MTP (mlx-lm PR 990 / PR 15 monkey-patch)
    if "mtp_enabled" in sent:
        new_mtp_enabled = False if is_diffusion_model else bool(request.mtp_enabled)
        if new_mtp_enabled:
            # Compatibility check: the model needs MTP heads in config.json AND
            # the model_type must be one PR 990 / PR 15 covers AND the weight
            # files must actually contain MTP tensors (mtp.* or the native
            # nextn layers). The last check is the one that catches
            # mlx-community converted weights where the default sanitize
            # path stripped the MTP heads.
            import json
            from pathlib import Path

            from ..utils.model_loading import (
                _checkpoint_has_mtp_weights,
                _is_mtp_compatible,
            )

            cfg_path = Path(entry.model_path) / "config.json"
            if not cfg_path.exists():
                raise HTTPException(
                    status_code=400,
                    detail=(
                        f"MTP enabled but config.json missing at {cfg_path}; "
                        "cannot verify MTP compatibility."
                    ),
                )
            try:
                cfg = json.loads(cfg_path.read_text())
            except Exception as e:
                raise HTTPException(
                    status_code=400,
                    detail=f"MTP enabled but failed to read model config: {e}",
                )
            model_type = cfg.get("model_type")
            if not _is_mtp_compatible(cfg, model_type):
                raise HTTPException(
                    status_code=400,
                    detail=(
                        f"Model is not MTP-compatible (model_type={model_type!r}, "
                        f"mtp_num_hidden_layers={cfg.get('mtp_num_hidden_layers', 0)}). "
                        "Lightning MTP requires a Qwen3.5/3.6, DeepSeek-V4 or "
                        "GLM-5.2 checkpoint with MTP heads."
                    ),
                )
            if not _checkpoint_has_mtp_weights(entry.model_path):
                raise HTTPException(
                    status_code=400,
                    detail=(
                        "Config declares MTP layers but the weight files contain "
                        "neither mtp.* tensors nor native nextn layers. Re-convert "
                        "from HF with a converter that preserves MTP weights. The "
                        "default mlx-lm sanitize() path strips them."
                    ),
                )
            # Mutual exclusion with DFlash — ModelSettings.__post_init__
            # also enforces this, but we surface a clearer error here.
            dflash_after = (
                bool(request.dflash_enabled)
                if "dflash_enabled" in sent
                else current_settings.dflash_enabled
            )
            if dflash_after:
                raise HTTPException(
                    status_code=400,
                    detail="MTP and DFlash cannot both be enabled; choose one speculative-decoding path.",
                )
        current_settings.mtp_enabled = new_mtp_enabled

    # VLM MTP (mlx-vlm f96138e+, gemma4_assistant drafter)
    if "vlm_mtp_enabled" in sent:
        new_vlm_mtp = False if is_diffusion_model else bool(request.vlm_mtp_enabled)
        if new_vlm_mtp:
            drafter_after = (
                request.vlm_mtp_draft_model
                if "vlm_mtp_draft_model" in sent
                else current_settings.vlm_mtp_draft_model
            )
            if not drafter_after:
                raise HTTPException(
                    status_code=400,
                    detail=(
                        "vlm_mtp_enabled requires vlm_mtp_draft_model "
                        "(path to a gemma4_assistant drafter, "
                        "e.g. 'gemma-4-26B-A4B-it-assistant')."
                    ),
                )
            # Mutex enforced again at ModelSettings.__post_init__ for
            # last-mile safety, but surface a clearer error here.
            for other_field, other_label in (
                ("dflash_enabled", "DFlash"),
                ("specprefill_enabled", "SpecPrefill"),
                ("mtp_enabled", "MTP"),
                ("turboquant_kv_enabled", "TurboQuant KV"),
            ):
                other_after = (
                    bool(getattr(request, other_field))
                    if other_field in sent
                    else getattr(current_settings, other_field)
                )
                if other_after:
                    raise HTTPException(
                        status_code=400,
                        detail=(
                            f"vlm_mtp_enabled and {other_label} cannot both be "
                            "enabled; choose one speculative-decoding path."
                        ),
                    )
        current_settings.vlm_mtp_enabled = new_vlm_mtp
    if "vlm_mtp_draft_model" in sent:
        current_settings.vlm_mtp_draft_model = request.vlm_mtp_draft_model or None
    if "vlm_mtp_draft_block_size" in sent:
        current_settings.vlm_mtp_draft_block_size = request.vlm_mtp_draft_block_size

    if "reasoning_parser" in sent:
        current_settings.reasoning_parser = request.reasoning_parser or None
    if "guided_grammar_enabled" in sent:
        current_settings.guided_grammar_enabled = (
            request.guided_grammar_enabled or False
        )
    if "guided_grammar" in sent:
        grammar = request.guided_grammar.strip() if request.guided_grammar else None
        current_settings.guided_grammar = grammar or None
    if request.is_pinned is not None:
        current_settings.is_pinned = request.is_pinned
        # Also update the engine pool entry
        entry.is_pinned = request.is_pinned
    if request.is_default is not None:
        current_settings.is_default = request.is_default
        # Update server_state.default_model if setting as default
        if request.is_default and server_state:
            server_state.default_model = model_id
    if request.is_hidden is not None:
        current_settings.is_hidden = request.is_hidden
    if request.is_favorite is not None:
        current_settings.is_favorite = request.is_favorite
    if "trust_remote_code" in sent:
        current_settings.trust_remote_code = bool(request.trust_remote_code)

    if is_diffusion_model:
        _sanitize_diffusion_model_settings(current_settings)

    # If an active profile was set, clear it when the user's save diverges
    # from the profile's stored values.  Only compare fields present in
    # both the profile and the current settings — new fields in the model
    # settings that the profile doesn't have are silently merged in, and
    # removed fields (no longer in the profile) are skipped.
    if current_settings.active_profile_name:
        profile = settings_manager.get_profile(
            model_id, current_settings.active_profile_name
        )
        if profile is None:
            current_settings.active_profile_name = None
        else:
            profile_settings = profile.get("settings", {}) or {}
            candidate = current_settings.to_dict()
            diverged = False
            for key, expected in profile_settings.items():
                # Profile None means "unconstrained" — candidate.to_dict()
                # drops None, so treat profile None as no constraint to
                # keep the comparison symmetric.
                if expected is None:
                    continue
                if key not in candidate:
                    diverged = True
                    break
                if candidate[key] != expected:
                    diverged = True
                    break
            if diverged:
                current_settings.active_profile_name = None
            else:
                new_fields = {
                    k: v
                    for k, v in candidate.items()
                    if k not in profile_settings and k not in EXCLUDED_FROM_PROFILES
                }
                if new_fields:
                    profile_settings.update(new_fields)
                    profile["settings"] = profile_settings
                    settings_manager.update_profile(
                        model_id,
                        current_settings.active_profile_name,
                        settings=profile_settings,
                    )

    # Persist settings
    settings_manager.set_settings(model_id, current_settings)

    # Auto-unload (and re-load if pinned) when a setting that only takes
    # effect at engine construction time is changed on a loaded model.
    requires_reload = entry.engine is not None and (
        ("model_type_override" in sent and entry.engine_type != prev_engine_type)
        or "index_cache_freq" in sent
        or "dflash_enabled" in sent
        or "dflash_draft_model" in sent
        or "dflash_draft_quant_enabled" in sent
        or "dflash_draft_quant_weight_bits" in sent
        or "dflash_draft_quant_activation_bits" in sent
        or "dflash_draft_quant_group_size" in sent
        or "dflash_max_ctx" in sent
        or "dflash_in_memory_cache" in sent
        or "dflash_in_memory_cache_max_entries" in sent
        or "dflash_in_memory_cache_max_bytes" in sent
        or "dflash_ssd_cache" in sent
        or "dflash_ssd_cache_max_bytes" in sent
        # trust_remote_code is plumbed at model load time; toggling it on
        # an already-loaded engine has no effect until reload.
        or "trust_remote_code" in sent
    )
    auto_unloaded = False
    auto_reloaded = False
    if requires_reload:
        was_pinned = entry.is_pinned
        try:
            logger.info(
                f"Settings changed for loaded model {model_id}, auto-unloading."
            )
            await engine_pool._unload_engine(model_id)
            auto_unloaded = True
        except Exception as e:
            logger.warning(f"Auto-unload failed for {model_id}: {e}")
        if auto_unloaded and was_pinned:
            try:
                await engine_pool._load_engine(model_id)
                auto_reloaded = True
                logger.info(f"Auto-reloaded pinned model {model_id} with new settings.")
            except Exception as e:
                logger.warning(f"Auto-reload failed for pinned model {model_id}: {e}")

    return {
        "success": True,
        "model_id": model_id,
        "settings": current_settings.to_dict(),
        "model_type": entry.model_type,
        "engine_type": entry.engine_type,
        "requires_reload": requires_reload,
        "auto_unloaded": auto_unloaded,
        "auto_reloaded": auto_reloaded,
    }


# =============================================================================
# Profile & Template endpoints
# =============================================================================


def _require_settings_manager():
    mgr = _get_settings_manager()
    if mgr is None:
        raise HTTPException(status_code=503, detail="Server not initialized")
    return mgr


def _require_model(model_id: str):
    pool = _get_engine_pool()
    if pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")
    entry = pool.get_entry(model_id)
    if entry is None:
        raise HTTPException(status_code=404, detail=f"Model not found: {model_id}")
    return entry


def _model_aliases(
    settings_manager, *, exclude_model_id: str | None = None
) -> dict[str, str]:
    return {
        ms.model_alias: mid
        for mid, ms in settings_manager.get_all_settings().items()
        if mid != exclude_model_id and ms.model_alias
    }


def _raise_if_profile_id_conflicts_model_id(
    candidate_id: str,
    *,
    model_id: str,
    engine_pool,
):
    for existing_id in engine_pool.get_model_ids():
        if existing_id != model_id and existing_id == candidate_id:
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Exposed profile model ID '{candidate_id}' conflicts with "
                    f"model directory name '{existing_id}'"
                ),
            )


def _raise_if_alias_conflicts_exposed_profiles(
    *,
    alias_value: str,
    model_id: str,
    settings_manager,
    engine_pool,
):
    exposed_ids = settings_manager.get_exposed_profile_model_ids()
    if alias_value in exposed_ids:
        raise HTTPException(
            status_code=400,
            detail=f"Alias '{alias_value}' conflicts with an exposed profile model ID",
        )

    aliases = _model_aliases(settings_manager, exclude_model_id=model_id)
    for profile in settings_manager.list_profiles(model_id):
        if not profile.get("expose_as_model"):
            continue
        api_name = profile.get("api_name") or profile["name"]
        candidate_id = f"{alias_value}:{api_name}"
        _raise_if_profile_id_conflicts_model_id(
            candidate_id,
            model_id=model_id,
            engine_pool=engine_pool,
        )
        if candidate_id in aliases:
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Alias '{alias_value}' would expose profile model ID "
                    f"'{candidate_id}', which conflicts with model alias "
                    f"for '{aliases[candidate_id]}'"
                ),
            )
        other_exposed_ids = settings_manager.get_exposed_profile_model_ids(
            exclude_model_id=model_id,
            exclude_profile_name=profile["name"],
        )
        if candidate_id in other_exposed_ids:
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Alias '{alias_value}' would expose duplicate profile "
                    f"model ID '{candidate_id}'"
                ),
            )


@router.get("/api/models/{model_id}/profiles")
async def list_model_profiles(
    model_id: str,
    is_admin: bool = Depends(require_admin),
):
    mgr = _require_settings_manager()
    _require_model(model_id)
    return {"profiles": mgr.list_profiles(model_id)}


@router.post("/api/models/{model_id}/profiles")
async def create_model_profile(
    model_id: str,
    request: CreateProfileRequest,
    is_admin: bool = Depends(require_admin),
):
    from ..model_profiles import InvalidProfileNameError, filter_universal_fields

    mgr = _require_settings_manager()
    _require_model(model_id)
    engine_pool = _get_engine_pool()
    try:
        profile = mgr.save_profile(
            model_id=model_id,
            name=request.name,
            display_name=request.display_name,
            description=request.description,
            settings=request.settings or {},
            source_template=request.source_template,
            expose_as_model=request.expose_as_model,
            api_name=request.api_name,
            reserved_model_ids=(
                set(engine_pool.get_model_ids()) if engine_pool is not None else None
            ),
        )
    except InvalidProfileNameError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=409, detail=str(e))

    if request.also_save_as_template:
        try:
            mgr.upsert_template(
                name=request.name,
                display_name=request.display_name,
                description=request.description,
                settings=filter_universal_fields(request.settings or {}),
            )
        except InvalidProfileNameError as e:
            raise HTTPException(status_code=400, detail=str(e))
    return {"profile": profile}


@router.put("/api/models/{model_id}/profiles/{name}")
async def update_model_profile(
    model_id: str,
    name: str,
    request: UpdateProfileRequest,
    is_admin: bool = Depends(require_admin),
):
    from ..model_profiles import InvalidProfileNameError, filter_universal_fields

    mgr = _require_settings_manager()
    _require_model(model_id)
    engine_pool = _get_engine_pool()
    try:
        updated = mgr.update_profile(
            model_id=model_id,
            name=name,
            new_name=request.new_name,
            display_name=request.display_name,
            description=request.description,
            settings=request.settings,
            source_template=request.source_template,
            expose_as_model=request.expose_as_model,
            api_name=request.api_name,
            reserved_model_ids=(
                set(engine_pool.get_model_ids()) if engine_pool is not None else None
            ),
        )
    except InvalidProfileNameError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=409, detail=str(e))
    if updated is None:
        raise HTTPException(status_code=404, detail=f"Profile not found: {name}")

    if request.also_save_as_template and request.settings is not None:
        try:
            mgr.upsert_template(
                name=updated["name"],
                display_name=updated["display_name"],
                description=updated.get("description"),
                settings=filter_universal_fields(request.settings),
            )
        except InvalidProfileNameError as e:
            raise HTTPException(status_code=400, detail=str(e))
    return {"profile": updated}


@router.delete("/api/models/{model_id}/profiles/{name}")
async def delete_model_profile(
    model_id: str,
    name: str,
    is_admin: bool = Depends(require_admin),
):
    mgr = _require_settings_manager()
    _require_model(model_id)
    if not mgr.delete_profile(model_id, name):
        raise HTTPException(status_code=404, detail=f"Profile not found: {name}")
    return {"deleted": True, "name": name}


@router.post("/api/models/{model_id}/profiles/{name}/apply")
async def apply_model_profile(
    model_id: str,
    name: str,
    is_admin: bool = Depends(require_admin),
):
    mgr = _require_settings_manager()
    entry = _require_model(model_id)
    is_diffusion_model = _entry_is_diffusion_model(entry)
    sanitizer = _sanitize_diffusion_settings_dict if is_diffusion_model else None
    applied = mgr.apply_profile(model_id, name, settings_sanitizer=sanitizer)
    if applied is None:
        raise HTTPException(status_code=404, detail=f"Profile not found: {name}")
    if is_diffusion_model:
        _sanitize_diffusion_model_settings(applied)
        mgr.set_settings(model_id, applied)
    return {"model_id": model_id, "settings": applied.to_dict()}


@router.get("/api/profile-fields")
async def get_profile_fields(is_admin: bool = Depends(require_admin)):
    from ..model_profiles import (
        MODEL_SPECIFIC_PROFILE_FIELDS,
        UNIVERSAL_PROFILE_FIELDS,
    )

    return {
        "universal": list(UNIVERSAL_PROFILE_FIELDS),
        "model_specific": list(MODEL_SPECIFIC_PROFILE_FIELDS),
    }


@router.get("/api/profile-templates")
async def list_templates(is_admin: bool = Depends(require_admin)):
    mgr = _require_settings_manager()
    return {"templates": mgr.list_templates()}


@router.post("/api/profile-templates")
async def create_template(
    request: CreateTemplateRequest,
    is_admin: bool = Depends(require_admin),
):
    from ..model_profiles import InvalidProfileNameError

    mgr = _require_settings_manager()
    try:
        tmpl = mgr.save_template(
            name=request.name,
            display_name=request.display_name,
            description=request.description,
            settings=request.settings or {},
        )
    except InvalidProfileNameError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=409, detail=str(e))
    return {"template": tmpl}


@router.put("/api/profile-templates/{name}")
async def update_template(
    name: str,
    request: UpdateTemplateRequest,
    is_admin: bool = Depends(require_admin),
):
    from ..model_profiles import InvalidProfileNameError

    mgr = _require_settings_manager()
    try:
        updated = mgr.update_template(
            name=name,
            new_name=request.new_name,
            display_name=request.display_name,
            description=request.description,
            settings=request.settings,
        )
    except InvalidProfileNameError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ValueError as e:
        raise HTTPException(status_code=409, detail=str(e))
    if updated is None:
        raise HTTPException(status_code=404, detail=f"Template not found: {name}")
    return {"template": updated}


@router.delete("/api/profile-templates/{name}")
async def delete_template(
    name: str,
    is_admin: bool = Depends(require_admin),
):
    mgr = _require_settings_manager()
    if not mgr.delete_template(name):
        raise HTTPException(status_code=404, detail=f"Template not found: {name}")
    return {"deleted": True, "name": name}


# =============================================================================
# Preset refresh (proxy to omlx.ai to avoid CORS)
# =============================================================================


@router.post("/api/presets/refresh")
async def refresh_presets(is_admin: bool = Depends(require_admin)):
    """Fetch the latest preset bundle from omlx.ai and return it.

    The client uses this instead of fetching omlx.ai directly so we do not
    depend on CORS headers on the remote host. Any failure is surfaced as 502
    so the client can silently fall back to the bundled presets.
    """
    try:
        resp = await asyncio.to_thread(
            requests.get,
            PRESET_REMOTE_URL,
            timeout=10,
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Fetch failed: {e}")
    if resp.status_code != 200:
        raise HTTPException(
            status_code=502,
            detail=f"Remote returned {resp.status_code}",
        )
    try:
        return resp.json()
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Invalid JSON: {e}")


@router.get("/api/models/{model_id}/generation_config")
async def get_generation_config(
    model_id: str,
    is_admin: bool = Depends(require_admin),
):
    """
    Read model config files and return recommended defaults.

    Reads generation_config.json for sampling parameters and config.json
    for max_context_window (max_position_embeddings).

    Args:
        model_id: The model identifier.

    Returns:
        JSON with recommended parameters from the model's config files.

    Raises:
        HTTPException: 404 if model not found or no config files exist.
    """
    import json as json_module

    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    entry = engine_pool.get_entry(model_id)
    if entry is None:
        raise HTTPException(status_code=404, detail=f"Model not found: {model_id}")

    model_path = Path(entry.model_path)
    result = {}

    # Read generation_config.json for sampling parameters
    gen_config_path = model_path / "generation_config.json"
    if gen_config_path.exists():
        try:
            with open(gen_config_path, encoding="utf-8") as f:
                gen_config = json_module.load(f)

            # Temperature: if do_sample is false, effective temperature is 0
            do_sample = gen_config.get("do_sample", True)
            if "temperature" in gen_config:
                result["temperature"] = (
                    0.0 if not do_sample else gen_config["temperature"]
                )

            if "top_p" in gen_config:
                result["top_p"] = gen_config["top_p"]

            if "top_k" in gen_config:
                result["top_k"] = gen_config["top_k"]

            if "repetition_penalty" in gen_config:
                result["repetition_penalty"] = gen_config["repetition_penalty"]

        except (json_module.JSONDecodeError, OSError) as e:
            logger.warning(
                f"Failed to parse generation_config.json for {model_id}: {e}"
            )

    # Read config.json for max_position_embeddings → max_context_window
    config_path = model_path / "config.json"
    if config_path.exists():
        try:
            with open(config_path, encoding="utf-8") as f:
                model_config = json_module.load(f)

            max_pos = (
                model_config.get("max_position_embeddings")
                or model_config.get("max_seq_len")
                or model_config.get("seq_length")
                or model_config.get("n_positions")
            )

            # Nested config fallback (VLM, MoE models like Qwen3.5, GLM-4V)
            if not max_pos:
                text_config = model_config.get("text_config", {})
                if isinstance(text_config, dict):
                    max_pos = (
                        text_config.get("max_position_embeddings")
                        or text_config.get("max_seq_len")
                        or text_config.get("seq_length")
                        or text_config.get("n_positions")
                    )

            if max_pos and isinstance(max_pos, int):
                result["max_context_window"] = max_pos

        except (json_module.JSONDecodeError, OSError) as e:
            logger.warning(f"Failed to parse config.json for {model_id}: {e}")

    if not result:
        raise HTTPException(
            status_code=404,
            detail=f"No config files with defaults found for {model_id}",
        )

    return result


# =============================================================================
# Global Settings API Routes
# =============================================================================


@router.get("/api/server-info")
async def get_server_info(is_admin: bool = Depends(require_admin)):
    """Return server connectivity metadata for the dashboard.

    Provides the configured host, port, and the list of user-facing
    aliases (hostnames/IPs) that the dashboard can use to render
    selectable API URL hints.

    Returns:
        JSON object with ``host``, ``port``, and ``aliases``.

    Raises:
        HTTPException: 401 if not authenticated, 503 if server not initialized.
    """
    from ..utils.network import detect_server_aliases

    global_settings = _get_global_settings()
    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    configured = list(global_settings.server.server_aliases)
    if configured:
        aliases = configured
    else:
        # Fall back to live detection if persisted list is empty.
        aliases = detect_server_aliases(host=global_settings.server.host)

    return {
        "host": global_settings.server.host,
        "port": global_settings.server.port,
        "aliases": aliases,
    }


def _schedule_self_terminate(delay: float = 0.5) -> None:
    """Schedule ``os.kill(getpid(), SIGTERM)`` on the running loop.

    Extracted from the restart handler so tests can patch this seam
    instead of mocking ``asyncio.get_running_loop`` globally (which
    interferes with FastAPI's TestClient portal).
    """
    pid = os.getpid()

    def _kill() -> None:
        try:
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            # Already exited (e.g. concurrent SIGTERM) — nothing to do.
            pass
        except Exception:  # pragma: no cover — best-effort signal.
            logger.exception("Failed to self-terminate for restart")

    asyncio.get_running_loop().call_later(delay, _kill)


@router.post("/api/server/restart")
async def restart_server(is_admin: bool = Depends(require_admin)):
    """Trigger a server restart via the menubar supervisor.

    The handler does not perform the restart itself — it returns 202 and
    schedules ``os.kill(os.getpid(), SIGTERM)`` 500ms after the response
    is queued. The menubar app's ``ServerManager._health_check_loop``
    detects the process exit and respawns the server with a short
    backoff (~5s).

    Gated by the ``OMLX_SUPERVISED`` environment variable so plain
    ``omlx serve`` (no supervisor) returns 503 rather than killing the
    server with no respawn path.
    """
    supervisor = os.environ.get("OMLX_SUPERVISED")
    if not supervisor:
        raise HTTPException(
            status_code=503,
            detail=(
                "Server is not running under a supervisor that can "
                "respawn it. Restart unavailable — use the menu bar "
                "app's Restart, or restart from your shell."
            ),
        )

    _schedule_self_terminate(0.5)
    logger.warning("Server restart requested (supervisor=%s)", supervisor)

    # 5s backoff in ServerManager + ~1-2s startup = ~7s downtime budget.
    return JSONResponse(
        status_code=202,
        content={
            "status": "restarting",
            "supervisor": supervisor,
            "expected_downtime_seconds": 7,
        },
    )


@router.get("/api/global-settings")
async def get_global_settings(is_admin: bool = Depends(require_admin)):
    """
    Get current global server settings.

    Returns the full global settings including server, model, scheduler,
    cache, and MCP configurations.

    Returns:
        JSON object with global settings.

    Raises:
        HTTPException: 401 if not authenticated, 503 if server not initialized.
    """
    global_settings = _get_global_settings()

    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Get system memory info for auto calculation
    memory_info = get_system_memory_info()

    # Get SSD disk info for cache directory
    cache_dir = global_settings.cache.ssd_cache_dir or str(
        global_settings.cache.get_ssd_cache_dir(global_settings.base_path)
    )
    disk_info = get_ssd_disk_info(cache_dir)

    return {
        "base_path": str(global_settings.base_path),
        "server": {
            "host": global_settings.server.host,
            "port": global_settings.server.port,
            "log_level": global_settings.server.log_level,
            "server_aliases": list(global_settings.server.server_aliases),
            "sse_keepalive_mode": global_settings.server.sse_keepalive_mode,
            "auto_start_on_launch": global_settings.server.auto_start_on_launch,
            "burst_decode_mode": global_settings.server.burst_decode_mode,
            "preserve_mid_system_cache": getattr(
                global_settings.server,
                "preserve_mid_system_cache",
                True,
            ),
        },
        "model": {
            "model_dirs": [
                str(d)
                for d in global_settings.model.get_model_dirs(global_settings.base_path)
            ],
            "model_dir": str(
                global_settings.model.get_model_dir(global_settings.base_path)
            ),
            "effective_model_dirs": [
                str(d) for d in global_settings.get_effective_model_dirs()
            ],
            "model_fallback": global_settings.model.model_fallback,
            "hide_helper_models": global_settings.model.hide_helper_models,
        },
        "memory": {
            "prefill_memory_guard": global_settings.memory.prefill_memory_guard,
            "memory_guard_tier": global_settings.memory.memory_guard_tier,
            "memory_guard_custom_ceiling_gb": global_settings.memory.memory_guard_custom_ceiling_gb,
        },
        "scheduler": {
            "max_concurrent_requests": global_settings.scheduler.max_concurrent_requests,
            "embedding_batch_size": global_settings.scheduler.embedding_batch_size,
            "chunked_prefill": global_settings.scheduler.chunked_prefill,
            "prefill_priority": global_settings.scheduler.prefill_priority,
        },
        "cache": {
            "enabled": global_settings.cache.enabled,
            "ssd_cache_dir": cache_dir,
            # Resolve "auto" to actual value (10% of SSD capacity)
            "ssd_cache_max_size": _format_cache_size(
                global_settings.cache.get_ssd_cache_max_size_bytes(
                    global_settings.base_path
                )
            ),
            "hot_cache_only": global_settings.cache.hot_cache_only,
            "hot_cache_max_size": global_settings.cache.hot_cache_max_size,
            "initial_cache_blocks": global_settings.cache.initial_cache_blocks,
        },
        "mcp": {
            "config_path": global_settings.mcp.config_path,
        },
        "huggingface": {
            "endpoint": global_settings.huggingface.endpoint,
            "hf_cache_enabled": global_settings.huggingface.hf_cache_enabled,
            "hf_cache_path": str(global_settings.get_hf_cache_dir()),
        },
        "modelscope": {
            "endpoint": global_settings.modelscope.endpoint,
        },
        "network": {
            "http_proxy": global_settings.network.http_proxy,
            "https_proxy": global_settings.network.https_proxy,
            "no_proxy": global_settings.network.no_proxy,
            "ca_bundle": global_settings.network.ca_bundle,
        },
        "sampling": {
            "max_context_window": global_settings.sampling.max_context_window,
            "max_context_window_policy": (
                global_settings.sampling.max_context_window_policy
            ),
            "max_tokens": global_settings.sampling.max_tokens,
            "temperature": global_settings.sampling.temperature,
            "top_p": global_settings.sampling.top_p,
            "top_k": global_settings.sampling.top_k,
            "repetition_penalty": global_settings.sampling.repetition_penalty,
        },
        "auth": {
            "api_key_set": bool(global_settings.auth.api_key),
            "api_key": global_settings.auth.api_key or "",
            "skip_api_key_verification": global_settings.auth.skip_api_key_verification,
            "sub_keys": [sk.to_dict() for sk in global_settings.auth.sub_keys],
        },
        "claude_code": {
            "mode": global_settings.claude_code.mode,
            "opus_model": global_settings.claude_code.opus_model,
            "sonnet_model": global_settings.claude_code.sonnet_model,
            "haiku_model": global_settings.claude_code.haiku_model,
        },
        "integrations": {
            "codex_model": global_settings.integrations.codex_model,
            "opencode_model": global_settings.integrations.opencode_model,
            "openclaw_model": global_settings.integrations.openclaw_model,
            "hermes_model": global_settings.integrations.hermes_model,
            "pi_model": global_settings.integrations.pi_model,
            "copilot_model": global_settings.integrations.copilot_model,
            "openclaw_tools_profile": global_settings.integrations.openclaw_tools_profile,
            "markitdown_enabled": global_settings.integrations.markitdown_enabled,
            "markitdown_expose_model": global_settings.integrations.markitdown_expose_model,
            "markitdown_max_file_size_mb": global_settings.integrations.markitdown_max_file_size_mb,
            "markitdown_max_files_per_request": global_settings.integrations.markitdown_max_files_per_request,
            "markitdown_pdf_processing_engine": global_settings.integrations.markitdown_pdf_processing_engine,
        },
        "system": {
            "total_memory_bytes": memory_info["total_bytes"],
            "total_memory": memory_info["total_formatted"],
            "auto_model_memory": memory_info["auto_limit_formatted"],
            "available_memory_bytes": memory_info["available_bytes"],
            "omlx_phys_footprint_bytes": memory_info["omlx_phys_footprint_bytes"],
            "free_memory_bytes": memory_info["free_memory_bytes"],
            "inactive_memory_bytes": memory_info["inactive_memory_bytes"],
            "active_memory_bytes": memory_info["active_memory_bytes"],
            "iogpu_wired_limit_bytes": memory_info["iogpu_wired_limit_bytes"],
            "omlx_wired_limit_request_bytes": memory_info[
                "omlx_wired_limit_request_bytes"
            ],
            "ssd_total_bytes": disk_info["total_bytes"],
            "ssd_total": disk_info["total_formatted"],
        },
        "ui": {
            "language": global_settings.ui.language,
        },
        "idle_timeout": {
            "idle_timeout_seconds": global_settings.idle_timeout.idle_timeout_seconds,
        },
    }


@router.post("/api/global-settings")
async def update_global_settings(
    request: GlobalSettingsRequest,
    is_admin: bool = Depends(require_admin),
):
    """
    Update global server settings.

    Updates are persisted to the global settings file. Some settings
    (log_level, model_dir, memory_guard_tier, cache) are applied immediately,
    while others (host, port, scheduler, mcp) require server restart.

    Args:
        request: GlobalSettingsRequest with the new settings.

    Returns:
        JSON response with success status, message, and list of runtime-applied settings.

    Raises:
        HTTPException: 401 if not authenticated, 503 if server not initialized,
                      400 if validation fails.
    """
    global_settings = _get_global_settings()

    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Track which settings were applied at runtime
    runtime_applied: list[str] = []
    pending_embedding_batch_size: int | None = None
    previous_embedding_batch_size: int | None = None

    # Apply server settings
    if request.host is not None:
        from ..utils.network import is_valid_bind_host

        parts = [h.strip() for h in request.host.split(",") if h.strip()]
        if not parts:
            raise HTTPException(status_code=400, detail="Host cannot be empty")
        for part in parts:
            if not is_valid_bind_host(part):
                raise HTTPException(
                    status_code=400,
                    detail=f"Invalid host: {part!r} (must be a hostname or IP address)",
                )
        global_settings.server.host = request.host
    if request.port is not None:
        global_settings.server.port = request.port
    if request.log_level is not None:
        global_settings.server.log_level = request.log_level
        # Apply log level at runtime
        _apply_log_level_runtime(request.log_level)
        runtime_applied.append("log_level")
    if request.sse_keepalive_mode is not None:
        valid_modes = {"chunk", "comment", "off"}
        if request.sse_keepalive_mode not in valid_modes:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid sse_keepalive_mode: {request.sse_keepalive_mode} "
                f"(must be one of {sorted(valid_modes)})",
            )
        global_settings.server.sse_keepalive_mode = request.sse_keepalive_mode
        runtime_applied.append("sse_keepalive_mode")
    if request.burst_decode_mode is not None:
        if request.burst_decode_mode not in BURST_DECODE_MODES:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid burst_decode_mode: {request.burst_decode_mode} "
                f"(must be one of {sorted(BURST_DECODE_MODES)})",
            )
        mode = request.burst_decode_mode
        global_settings.server.burst_decode_mode = mode
        # Seed env so models loaded later pick up the mode without a restart.
        for _key, _value in burst_decode_env(mode).items():
            os.environ[_key] = _value
        # Hot-apply to every loaded engine. EngineConfig is a mutable dataclass
        # and its burst fields are read fresh each decode burst
        # (EngineCore._step_burst), so this takes effect on the next token.
        max_steps, single_s = BURST_DECODE_MODES[mode]
        from ..server import _server_state

        pool = _server_state.engine_pool
        if pool is not None:
            for _mid, entry in pool._entries.items():
                if entry is None or entry.engine is None:
                    continue
                async_core = getattr(entry.engine, "_engine", None)
                core = (
                    getattr(async_core, "engine", None)
                    if async_core is not None
                    else None
                )
                cfg = getattr(core, "config", None) if core is not None else None
                if cfg is not None and hasattr(cfg, "decode_burst_budget_single_s"):
                    cfg.decode_burst_max_steps = max_steps
                    cfg.decode_burst_budget_single_s = single_s
        runtime_applied.append("burst_decode_mode")
        logger.info(f"Burst Decode mode set to '{mode}'")
    if request.auto_start_on_launch is not None:
        global_settings.server.auto_start_on_launch = request.auto_start_on_launch
        runtime_applied.append("auto_start_on_launch")
    if request.preserve_mid_system_cache is not None:
        global_settings.server.preserve_mid_system_cache = (
            request.preserve_mid_system_cache
        )
        runtime_applied.append("preserve_mid_system_cache")

    if request.server_aliases is not None:
        from ..utils.network import is_valid_alias

        cleaned: list[str] = []
        seen: set[str] = set()
        for alias in request.server_aliases:
            if not isinstance(alias, str):
                raise HTTPException(
                    status_code=400,
                    detail="Invalid server alias: each alias must be a string",
                )
            value = alias.strip()
            if not value or value in seen:
                continue
            if not is_valid_alias(value):
                raise HTTPException(
                    status_code=400,
                    detail=f"Invalid server alias: {value!r} (must be a hostname or IP address)",
                )
            seen.add(value)
            cleaned.append(value)
        global_settings.server.server_aliases = cleaned
        runtime_applied.append("server_aliases")

    # Apply model settings
    new_dirs = None
    if request.model_dirs is not None:
        new_dirs = [d for d in request.model_dirs if d.strip()]
    elif request.model_dir is not None:
        new_dirs = [request.model_dir]

    if new_dirs is not None:
        old_dirs = global_settings.model.model_dirs
        if new_dirs != old_dirs:
            effective_dirs = [
                str(d) for d in global_settings.get_effective_model_dirs(new_dirs)
            ]
            success, msg = await _apply_model_dirs_runtime(effective_dirs)
            if success:
                global_settings.model.model_dirs = new_dirs
                global_settings.model.model_dir = new_dirs[0] if new_dirs else None
                runtime_applied.append("model_dirs")
                logger.info(msg)
            else:
                raise HTTPException(
                    status_code=400, detail=f"Failed to change model directories: {msg}"
                )

    if request.model_fallback is not None:
        global_settings.model.model_fallback = request.model_fallback
        runtime_applied.append("model_fallback")
    if request.hide_helper_models is not None:
        global_settings.model.hide_helper_models = request.hide_helper_models
        runtime_applied.append("hide_helper_models")

    # Apply memory guard tier + custom ceiling change (Live)
    if (
        request.memory_guard_tier is not None
        or request.memory_guard_custom_ceiling_gb is not None
    ):
        if request.memory_guard_tier is not None:
            global_settings.memory.memory_guard_tier = request.memory_guard_tier
        if request.memory_guard_custom_ceiling_gb is not None:
            global_settings.memory.memory_guard_custom_ceiling_gb = float(
                request.memory_guard_custom_ceiling_gb
            )
        try:
            success, msg = await _apply_memory_guard_tier_runtime(
                tier=request.memory_guard_tier,
                custom_ceiling_gb=request.memory_guard_custom_ceiling_gb,
            )
            if success:
                runtime_applied.append("memory_guard_tier")
                logger.info(msg)
            else:
                logger.warning(f"Failed to apply memory_guard_tier: {msg}")
        except Exception as e:
            logger.warning(f"Error applying memory_guard_tier: {e}")

    # Apply prefill memory guard setting (Live)
    if request.memory_prefill_memory_guard is not None:
        global_settings.memory.prefill_memory_guard = (
            request.memory_prefill_memory_guard
        )
        from ..server import _server_state

        if _server_state.process_memory_enforcer is not None:
            _server_state.process_memory_enforcer.prefill_memory_guard = (
                request.memory_prefill_memory_guard
            )
        runtime_applied.append("prefill_memory_guard")
        logger.info(
            f"Prefill memory guard "
            f"{'enabled' if request.memory_prefill_memory_guard else 'disabled'}"
        )

    # Apply scheduler settings (restart required)
    if request.max_concurrent_requests is not None:
        global_settings.scheduler.max_concurrent_requests = (
            request.max_concurrent_requests
        )

    # Apply embedding batch size setting (Live for loaded embedding engines)
    if request.embedding_batch_size is not None:
        if request.embedding_batch_size <= 0:
            raise HTTPException(
                status_code=400,
                detail="Invalid embedding_batch_size: must be > 0",
            )
        pending_embedding_batch_size = request.embedding_batch_size

    # Apply chunked prefill setting (Live)
    if request.chunked_prefill is not None:
        global_settings.scheduler.chunked_prefill = request.chunked_prefill
        from ..server import _server_state

        pool = _server_state.engine_pool
        if pool is not None:
            for mid, entry in pool._entries.items():
                if entry is None or entry.engine is None:
                    continue
                async_core = getattr(entry.engine, "_engine", None)
                core = (
                    getattr(async_core, "engine", None)
                    if async_core is not None
                    else None
                )
                scheduler = (
                    getattr(core, "scheduler", None) if core is not None else None
                )
                if scheduler is not None and hasattr(scheduler, "config"):
                    scheduler.config.chunked_prefill = request.chunked_prefill
        runtime_applied.append("chunked_prefill")
        logger.info(
            f"Chunked prefill {'enabled' if request.chunked_prefill else 'disabled'}"
        )

    # Apply prefill priority setting (Live)
    if request.prefill_priority is not None:
        value = request.prefill_priority.strip().lower()
        if value not in ("context", "speed"):
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Invalid prefill_priority: '{request.prefill_priority}' "
                    f"(must be 'context' or 'speed')"
                ),
            )
        global_settings.scheduler.prefill_priority = value
        from ..server import _server_state

        pool = _server_state.engine_pool
        if pool is not None:
            # Engines loaded from now on build their Scheduler from the
            # pool's stored config — without this, a bench/reload after the
            # toggle would silently revert to the boot-time mode.
            pool_config = getattr(pool, "_scheduler_config", None)
            if pool_config is not None:
                pool_config.prefill_speed_priority = value == "speed"
            for mid, entry in pool._entries.items():
                if entry is None or entry.engine is None:
                    continue
                async_core = getattr(entry.engine, "_engine", None)
                core = (
                    getattr(async_core, "engine", None)
                    if async_core is not None
                    else None
                )
                scheduler = (
                    getattr(core, "scheduler", None) if core is not None else None
                )
                if scheduler is not None:
                    scheduler._prefill_speed_priority = value == "speed"
                    if hasattr(scheduler, "config"):
                        scheduler.config.prefill_speed_priority = value == "speed"
        runtime_applied.append("prefill_priority")
        logger.info(f"Prefill priority set to '{value}'")

    if request.hot_cache_max_size is not None:
        try:
            _parse_hot_cache_max_size(request.hot_cache_max_size)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    # Apply cache settings
    cache_changed = False
    if request.cache_enabled is not None:
        global_settings.cache.enabled = request.cache_enabled
        cache_changed = True
    if request.ssd_cache_dir is not None:
        global_settings.cache.ssd_cache_dir = request.ssd_cache_dir
        cache_changed = True
    if request.ssd_cache_max_size is not None:
        global_settings.cache.ssd_cache_max_size = request.ssd_cache_max_size
        cache_changed = True
    if request.hot_cache_only is not None:
        global_settings.cache.hot_cache_only = request.hot_cache_only
    if request.hot_cache_max_size is not None:
        global_settings.cache.hot_cache_max_size = request.hot_cache_max_size
        cache_changed = True
    if request.initial_cache_blocks is not None:
        global_settings.cache.initial_cache_blocks = request.initial_cache_blocks

    if cache_changed:
        success, msg = await _apply_cache_settings_runtime(
            request.cache_enabled,
            request.ssd_cache_dir,
            request.ssd_cache_max_size,
            global_settings,
            hot_cache_max_size=request.hot_cache_max_size,
        )
        if success:
            runtime_applied.append("cache")
            logger.info(msg)
        else:
            logger.warning(f"Failed to apply cache settings runtime: {msg}")

    # Apply MCP settings (restart required)
    if request.mcp_config is not None:
        global_settings.mcp.config_path = (
            request.mcp_config if request.mcp_config else None
        )

    # Apply HuggingFace settings (Live - immediately applied via env var)
    if request.hf_endpoint is not None:
        global_settings.huggingface.endpoint = request.hf_endpoint
        if request.hf_endpoint:
            os.environ["HF_ENDPOINT"] = request.hf_endpoint
        elif "HF_ENDPOINT" in os.environ:
            del os.environ["HF_ENDPOINT"]
        runtime_applied.append("hf_endpoint")
        logger.info(
            f"HuggingFace endpoint updated to: " f"{request.hf_endpoint or '(default)'}"
        )
    if request.hf_cache_enabled is not None:
        if global_settings.huggingface.hf_cache_enabled != request.hf_cache_enabled:
            global_settings.huggingface.hf_cache_enabled = request.hf_cache_enabled
            effective_dirs = [
                str(d) for d in global_settings.get_effective_model_dirs()
            ]
            success, msg = await _apply_model_dirs_runtime(effective_dirs)
            if not success:
                raise HTTPException(
                    status_code=400,
                    detail=f"Failed to change HuggingFace cache discovery: {msg}",
                )
            runtime_applied.append("hf_cache_enabled")
            logger.info(msg)

    # Apply ModelScope settings (Live - immediately applied via env var)
    if request.ms_endpoint is not None:
        global_settings.modelscope.endpoint = request.ms_endpoint
        if request.ms_endpoint:
            os.environ["MODELSCOPE_DOMAIN"] = request.ms_endpoint
        elif "MODELSCOPE_DOMAIN" in os.environ:
            del os.environ["MODELSCOPE_DOMAIN"]
        runtime_applied.append("ms_endpoint")
        logger.info(
            f"ModelScope endpoint updated to: " f"{request.ms_endpoint or '(default)'}"
        )

    # Apply network settings (Live - immediately applied via env vars)
    network_changed = False
    if request.network_http_proxy is not None:
        global_settings.network.http_proxy = request.network_http_proxy
        if request.network_http_proxy:
            os.environ["HTTP_PROXY"] = request.network_http_proxy
            os.environ["http_proxy"] = request.network_http_proxy
        else:
            os.environ.pop("HTTP_PROXY", None)
            os.environ.pop("http_proxy", None)
        network_changed = True

    if request.network_https_proxy is not None:
        global_settings.network.https_proxy = request.network_https_proxy
        if request.network_https_proxy:
            os.environ["HTTPS_PROXY"] = request.network_https_proxy
            os.environ["https_proxy"] = request.network_https_proxy
        else:
            os.environ.pop("HTTPS_PROXY", None)
            os.environ.pop("https_proxy", None)
        network_changed = True

    if request.network_no_proxy is not None:
        global_settings.network.no_proxy = request.network_no_proxy
        if request.network_no_proxy:
            os.environ["NO_PROXY"] = request.network_no_proxy
            os.environ["no_proxy"] = request.network_no_proxy
        else:
            os.environ.pop("NO_PROXY", None)
            os.environ.pop("no_proxy", None)
        network_changed = True

    if request.network_ca_bundle is not None:
        global_settings.network.ca_bundle = request.network_ca_bundle
        if request.network_ca_bundle:
            os.environ["REQUESTS_CA_BUNDLE"] = request.network_ca_bundle
            os.environ["SSL_CERT_FILE"] = request.network_ca_bundle
        else:
            os.environ.pop("REQUESTS_CA_BUNDLE", None)
            os.environ.pop("SSL_CERT_FILE", None)
        network_changed = True

    if network_changed:
        runtime_applied.append("network")
        logger.info("Network settings updated")

    # Apply sampling settings (Live - immediately applied)
    sampling_changed = False
    if request.sampling_max_context_window is not None:
        global_settings.sampling.max_context_window = (
            request.sampling_max_context_window
        )
        sampling_changed = True
    if "sampling_max_context_window_policy" in request.model_fields_set:
        global_settings.sampling.max_context_window_policy = (
            request.sampling_max_context_window_policy
        )
        sampling_changed = True
    if request.sampling_max_tokens is not None:
        global_settings.sampling.max_tokens = request.sampling_max_tokens
        sampling_changed = True
    if request.sampling_temperature is not None:
        global_settings.sampling.temperature = request.sampling_temperature
        sampling_changed = True
    if request.sampling_top_p is not None:
        global_settings.sampling.top_p = request.sampling_top_p
        sampling_changed = True
    if request.sampling_top_k is not None:
        global_settings.sampling.top_k = request.sampling_top_k
        sampling_changed = True
    if request.sampling_repetition_penalty is not None:
        global_settings.sampling.repetition_penalty = (
            request.sampling_repetition_penalty
        )
        sampling_changed = True

    if sampling_changed:
        success, msg = _apply_sampling_settings_runtime(
            request.sampling_max_context_window,
            request.sampling_max_context_window_policy,
            "sampling_max_context_window_policy" in request.model_fields_set,
            request.sampling_max_tokens,
            request.sampling_temperature,
            request.sampling_top_p,
            request.sampling_top_k,
            request.sampling_repetition_penalty,
        )
        if success:
            runtime_applied.append("sampling")
            logger.info(msg)

    # Apply Claude Code settings (Live - immediately applied)
    claude_code_changed = False
    # mode: standard is-not-None check is correct — mode must never be null
    if request.claude_code_mode is not None:
        global_settings.claude_code.mode = request.claude_code_mode
        claude_code_changed = True
    # model fields: use model_fields_set to distinguish "field absent from POST body"
    # from "field explicitly sent as null" — null must clear the field to None.
    # DO NOT use `is not None` here: that would prevent clearing a model field to null.
    if "claude_code_opus_model" in request.model_fields_set:
        global_settings.claude_code.opus_model = request.claude_code_opus_model
        claude_code_changed = True
    if "claude_code_sonnet_model" in request.model_fields_set:
        global_settings.claude_code.sonnet_model = request.claude_code_sonnet_model
        claude_code_changed = True
    if "claude_code_haiku_model" in request.model_fields_set:
        global_settings.claude_code.haiku_model = request.claude_code_haiku_model
        claude_code_changed = True

    if claude_code_changed:
        runtime_applied.append("claude_code")
        logger.info(
            f"Claude Code settings updated: "
            f"mode={global_settings.claude_code.mode}, "
            f"opus={global_settings.claude_code.opus_model}, "
            f"sonnet={global_settings.claude_code.sonnet_model}, "
            f"haiku={global_settings.claude_code.haiku_model}"
        )

    # Apply integrations settings (Live - immediately applied)
    integrations_changed = False
    if "integrations_copilot_model" in request.model_fields_set:
        global_settings.integrations.copilot_model = request.integrations_copilot_model
        integrations_changed = True
    if "integrations_codex_model" in request.model_fields_set:
        global_settings.integrations.codex_model = request.integrations_codex_model
        integrations_changed = True
    if "integrations_opencode_model" in request.model_fields_set:
        global_settings.integrations.opencode_model = (
            request.integrations_opencode_model
        )
        integrations_changed = True
    if "integrations_openclaw_model" in request.model_fields_set:
        global_settings.integrations.openclaw_model = (
            request.integrations_openclaw_model
        )
        integrations_changed = True
    if "integrations_hermes_model" in request.model_fields_set:
        global_settings.integrations.hermes_model = request.integrations_hermes_model
        integrations_changed = True
    if "integrations_pi_model" in request.model_fields_set:
        global_settings.integrations.pi_model = request.integrations_pi_model
        integrations_changed = True
    if "integrations_openclaw_tools_profile" in request.model_fields_set:
        global_settings.integrations.openclaw_tools_profile = (
            request.integrations_openclaw_tools_profile
        )
        integrations_changed = True
    if "markitdown_enabled" in request.model_fields_set:
        global_settings.integrations.markitdown_enabled = bool(
            request.markitdown_enabled
        )
        integrations_changed = True
    if "markitdown_expose_model" in request.model_fields_set:
        global_settings.integrations.markitdown_expose_model = bool(
            request.markitdown_expose_model
        )
        integrations_changed = True
    if "markitdown_max_file_size_mb" in request.model_fields_set:
        if (
            request.markitdown_max_file_size_mb is None
            or request.markitdown_max_file_size_mb <= 0
        ):
            raise HTTPException(
                status_code=400,
                detail="markitdown_max_file_size_mb must be > 0",
            )
        global_settings.integrations.markitdown_max_file_size_mb = (
            request.markitdown_max_file_size_mb
        )
        integrations_changed = True
    if "markitdown_max_files_per_request" in request.model_fields_set:
        if (
            request.markitdown_max_files_per_request is None
            or request.markitdown_max_files_per_request <= 0
        ):
            raise HTTPException(
                status_code=400,
                detail="markitdown_max_files_per_request must be > 0",
            )
        global_settings.integrations.markitdown_max_files_per_request = (
            request.markitdown_max_files_per_request
        )
        integrations_changed = True
    if "markitdown_pdf_processing_engine" in request.model_fields_set:
        engine = (request.markitdown_pdf_processing_engine or "").strip()
        if not engine:
            raise HTTPException(
                status_code=400,
                detail="markitdown_pdf_processing_engine must not be empty",
            )
        global_settings.integrations.markitdown_pdf_processing_engine = engine
        integrations_changed = True

    if integrations_changed:
        runtime_applied.append("integrations")
        logger.info(
            f"Integration settings updated: "
            f"copilot={global_settings.integrations.copilot_model}, "
            f"codex={global_settings.integrations.codex_model}, "
            f"opencode={global_settings.integrations.opencode_model}, "
            f"openclaw={global_settings.integrations.openclaw_model}, "
            f"hermes={global_settings.integrations.hermes_model}, "
            f"pi={global_settings.integrations.pi_model}, "
            f"markitdown_enabled={global_settings.integrations.markitdown_enabled}, "
            f"markitdown_expose_model={global_settings.integrations.markitdown_expose_model}, "
            f"markitdown_pdf_processing_engine={global_settings.integrations.markitdown_pdf_processing_engine}"
        )

    # Apply UI settings
    if request.ui_language is not None:
        global_settings.ui.language = request.ui_language
        runtime_applied.append("ui_language")
        _refresh_i18n_globals()
        logger.info(f"UI language changed to: {request.ui_language}")

    # Apply idle timeout settings (Live)
    # Use model_fields_set to distinguish "explicitly sent as null" (disable)
    # from "not sent" (don't touch).
    if "idle_timeout_seconds" in request.model_fields_set:
        global_settings.idle_timeout.idle_timeout_seconds = request.idle_timeout_seconds
        runtime_applied.append("idle_timeout_seconds")
        if request.idle_timeout_seconds:
            logger.info(f"Idle timeout set to: {request.idle_timeout_seconds}s")
        else:
            logger.info("Idle timeout disabled")

    # Apply auth settings (API key change)
    if request.api_key is not None:
        from ..server import _server_state

        is_valid, error_msg = validate_api_key(request.api_key)
        if not is_valid:
            raise HTTPException(status_code=400, detail=error_msg)

        global_settings.auth.api_key = request.api_key
        _server_state.api_key = request.api_key
        runtime_applied.append("api_key")
        logger.info("API key updated via admin settings")

    if request.skip_api_key_verification is not None:
        global_settings.auth.skip_api_key_verification = (
            request.skip_api_key_verification
        )
        runtime_applied.append("skip_api_key_verification")

    if pending_embedding_batch_size is not None:
        previous_embedding_batch_size = global_settings.scheduler.embedding_batch_size
        global_settings.scheduler.embedding_batch_size = pending_embedding_batch_size

    # Validate settings
    errors = global_settings.validate()
    if errors:
        if previous_embedding_batch_size is not None:
            global_settings.scheduler.embedding_batch_size = (
                previous_embedding_batch_size
            )
        raise HTTPException(status_code=400, detail=errors)

    # Persist to file
    try:
        global_settings.save()
    except Exception as e:
        if previous_embedding_batch_size is not None:
            global_settings.scheduler.embedding_batch_size = (
                previous_embedding_batch_size
            )
        raise HTTPException(status_code=500, detail=f"Failed to save settings: {e}")

    if pending_embedding_batch_size is not None:
        from ..server import _server_state

        pool = _server_state.engine_pool
        if pool is not None:
            await pool.apply_embedding_batch_size(pending_embedding_batch_size)
        runtime_applied.append("embedding_batch_size")
        logger.info(f"Embedding batch size set to {pending_embedding_batch_size}")

    # Build response message
    message = "Settings saved successfully."

    return {
        "success": True,
        "message": message,
        "runtime_applied": runtime_applied,
    }


# =============================================================================
# Logs API Routes
# =============================================================================


def _tail_file(file_path: Path, num_lines: int) -> tuple[str, int]:
    """
    Read the last N lines of a file efficiently.

    Uses a deque to efficiently keep only the last N lines in memory.

    Args:
        file_path: Path to the log file.
        num_lines: Number of lines to return.

    Returns:
        Tuple of (content_string, total_line_count)
    """
    if not file_path.exists():
        return "", 0

    # Use deque for efficient tail operation
    lines = deque(maxlen=num_lines)
    total_lines = 0

    with open(file_path, encoding="utf-8", errors="replace") as f:
        for line in f:
            lines.append(line)
            total_lines += 1

    return "".join(lines), total_lines


def _get_available_log_files(log_dir: Path) -> list[str]:
    """
    Get list of available log files sorted by modification time.

    Args:
        log_dir: Directory containing log files.

    Returns:
        List of log file names, newest first.
    """
    if not log_dir.exists():
        return []

    files = []
    for f in log_dir.iterdir():
        # Match server.log and server.log.YYYY-MM-DD patterns
        if f.name.startswith("server") and (f.suffix == ".log" or ".log." in f.name):
            files.append(f.name)

    # Sort by modification time (newest first)
    files.sort(key=lambda x: (log_dir / x).stat().st_mtime, reverse=True)
    return files


@router.get("/api/logs")
async def get_logs(
    lines: int = 100,
    file: str | None = None,
    is_admin: bool = Depends(require_admin),
):
    """
    Get server logs.

    Returns the last N lines of the specified log file (or current log).
    Supports viewing historical rotated log files.

    Args:
        lines: Number of lines to return (default: 100, max: 10000).
        file: Optional specific log file name. If not specified, uses current log.

    Returns:
        JSON response with log content and metadata:
        - logs: The log content string
        - total_lines: Total number of lines in the file
        - log_file: Name of the log file being read
        - available_files: List of available log files

    Raises:
        HTTPException: 401 if not authenticated, 503 if server not initialized,
                      400 if invalid file name, 404 if log file not found.
    """
    global_settings = _get_global_settings()

    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    # Limit lines to prevent memory issues
    lines = min(max(1, lines), 10000)

    log_dir = global_settings.logging.get_log_dir(global_settings.base_path)

    # Get available log files
    available_files = _get_available_log_files(log_dir)

    # Determine which file to read
    if file:
        # Validate file name (prevent path traversal)
        if "/" in file or "\\" in file or ".." in file:
            raise HTTPException(status_code=400, detail="Invalid file name")
        log_file = log_dir / file
        if not log_file.exists():
            raise HTTPException(status_code=404, detail=f"Log file not found: {file}")
    else:
        # Default to current log file
        log_file = log_dir / "server.log"

    # Read log content
    if log_file.exists():
        content, total_lines = _tail_file(log_file, lines)
    else:
        content = ""
        total_lines = 0

    return {
        "logs": content,
        "total_lines": total_lines,
        "log_file": log_file.name,
        "available_files": available_files,
    }


# =============================================================================
# Stats API Routes
# =============================================================================


def _get_engine_info() -> dict:
    """Get commit SHA and GitHub URL for engine packages.

    Fallback chain:
    1. PEP 610 direct_url.json (pip install git+https://...)
    2. _engine_commits.json (generated by build.py for app bundle)
    3. Parse pyproject.toml at runtime (dev environment)
    """
    import importlib.metadata

    engines = {}
    packages = {
        "mlx-lm": "https://github.com/ml-explore/mlx-lm",
        "mlx-vlm": "https://github.com/Blaizzy/mlx-vlm",
        "mlx-embeddings": "https://github.com/Blaizzy/mlx-embeddings",
        "mlx-audio": "https://github.com/Blaizzy/mlx-audio",
    }

    fallback_commits = _load_fallback_commits(packages)

    for pkg_name, default_url in packages.items():
        info = {"name": pkg_name, "version": None, "commit": None, "url": None}
        try:
            dist = importlib.metadata.distribution(pkg_name)
            info["version"] = dist.version

            # Method 1: PEP 610 direct_url.json
            commit_info = _get_commit_from_direct_url(dist, default_url)
            if not commit_info:
                # Methods 2+3: _engine_commits.json or pyproject.toml
                commit_info = fallback_commits.get(pkg_name)

            if commit_info:
                info["commit"] = commit_info["commit"]
                info["url"] = commit_info["url"]
        except Exception:
            pass
        engines[pkg_name] = info

    return engines


def _get_commit_from_direct_url(dist, default_url: str) -> dict | None:
    """Extract commit SHA from PEP 610 direct_url.json."""
    import json

    try:
        direct_url_text = dist.read_text("direct_url.json")
        if direct_url_text:
            direct_url = json.loads(direct_url_text)
            vcs_info = direct_url.get("vcs_info", {})
            commit = vcs_info.get("commit_id")
            if commit:
                repo_url = direct_url.get("url", default_url).rstrip("/")
                if repo_url.endswith(".git"):
                    repo_url = repo_url[:-4]
                return {"commit": commit, "url": f"{repo_url}/commit/{commit}"}
    except Exception:
        pass
    return None


def _load_fallback_commits(packages: dict[str, str]) -> dict:
    """Load commit SHAs from fallback sources.

    Tries in order:
    1. _engine_commits.json (generated by build.py, lives in omlx package dir)
    2. pyproject.toml (dev environment, lives one level above package dir)
    """
    import json
    from pathlib import Path

    # This file is at omlx/admin/routes.py → package dir is omlx/
    pkg_dir = Path(__file__).resolve().parent.parent

    # Method 2: _engine_commits.json (written by build.py for app bundle)
    commits_file = pkg_dir / "_engine_commits.json"
    if commits_file.is_file():
        try:
            data = json.loads(commits_file.read_text())
            result = {}
            for pkg_name, entry in data.items():
                if isinstance(entry, dict) and "commit" in entry:
                    commit = entry["commit"]
                    repo_url = entry.get("url", packages.get(pkg_name, ""))
                    if "/commit/" not in repo_url:
                        repo_url = f"{repo_url}/commit/{commit}"
                    result[pkg_name] = {"commit": commit, "url": repo_url}
            if result:
                return result
        except Exception:
            pass

    # Method 3: Parse pyproject.toml (dev environment)
    pyproject = pkg_dir.parent / "pyproject.toml"
    if pyproject.is_file():
        try:
            return _parse_commits_from_pyproject(pyproject, packages)
        except Exception:
            pass

    return {}


def _parse_commits_from_pyproject(pyproject_path, packages: dict[str, str]) -> dict:
    """Extract commit SHAs from git+https:// URLs in pyproject.toml."""
    import re
    from pathlib import Path

    content = Path(pyproject_path).read_text()
    commits = {}
    # Match: "mlx-lm @ git+https://github.com/.../mlx-lm@<sha>"
    pattern = r'"(\S+)\s*@\s*git\+https://[^@"]+@([0-9a-f]{7,40})"'
    for match in re.finditer(pattern, content):
        pkg_name = match.group(1).strip().lower().split("[", 1)[0]
        sha = match.group(2)
        if pkg_name in packages:
            repo_url = packages[pkg_name]
            commits[pkg_name] = {
                "commit": sha,
                "url": f"{repo_url}/commit/{sha}",
            }
    return commits


def _build_runtime_cache_observability(
    global_settings,
    model_filter: str = "",
) -> dict:
    """Build runtime cache observability payload for dashboard.

    Includes the effective runtime paths and per-model SSD cache runtime stats
    from loaded schedulers, so users can verify real cache state without manual
    process inspection.
    """
    if global_settings is None:
        return {
            "base_path": "",
            "ssd_cache_dir": "",
            "response_state_dir": "",
            "models": [],
            "total_num_files": 0,
            "total_size_bytes": 0,
            "effective_block_sizes": [],
        }

    cache_dir = global_settings.cache.get_ssd_cache_dir(global_settings.base_path)
    cache_cfg = global_settings.cache
    try:
        cfg_disk_max = cache_cfg.get_ssd_cache_max_size_bytes(global_settings.base_path)
    except (ValueError, OSError, TypeError) as exc:
        logger.warning("Could not read SSD cache max size from config: %s", exc)
        cfg_disk_max = 0

    payload = {
        "base_path": str(global_settings.base_path),
        "ssd_cache_dir": str(cache_dir),
        "response_state_dir": str(cache_dir / "response-state"),
        "models": [],
        "total_num_files": 0,
        "total_size_bytes": 0,
        "effective_block_sizes": [],
        "disk_max_bytes": cfg_disk_max,
        "hot_cache_max_bytes": 0,
        "hot_cache_size_bytes": 0,
        "hot_cache_entries": 0,
    }

    engine_pool = _get_engine_pool()
    if engine_pool is None:
        return payload

    block_sizes = set()

    for model_info in engine_pool.get_status().get("models", []):
        model_id = model_info.get("id")
        if not model_id:
            continue
        if model_filter and model_id != model_filter:
            continue
        if not model_info.get("loaded"):
            continue

        entry = engine_pool._entries.get(model_id)
        if entry is None or entry.engine is None:
            continue

        async_core = getattr(entry.engine, "_engine", None)
        core = getattr(async_core, "engine", None) if async_core is not None else None
        scheduler = getattr(core, "scheduler", None) if core is not None else None
        if scheduler is None and async_core is None:
            # Engines without an AsyncEngineCore (DFlash) expose a scheduler
            # through their fallback engine once it is active.
            scheduler = getattr(entry.engine, "scheduler", None)

        runtime_stats = None
        if scheduler is not None and hasattr(scheduler, "get_ssd_cache_stats"):
            try:
                runtime_stats = scheduler.get_ssd_cache_stats()
            except Exception as exc:
                logger.warning(
                    "Failed to collect runtime cache stats for model '%s': %s",
                    model_id,
                    exc,
                )
                continue
        elif hasattr(entry.engine, "get_runtime_cache_stats"):
            # DFlash primary mode: the engine adapts its dflash-mlx runtime
            # cache (L1 in-memory + L2 snapshot dir) to the same shape.
            try:
                runtime_stats = entry.engine.get_runtime_cache_stats()
            except Exception as exc:  # noqa: BLE001
                logger.warning(
                    "Failed to collect runtime cache stats for model '%s': %s",
                    model_id,
                    exc,
                )
                continue

        if not runtime_stats:
            continue

        block_size = runtime_stats.get("block_size")
        indexed_blocks = runtime_stats.get("indexed_blocks")

        ssd_stats = runtime_stats.get("ssd_cache")
        if is_dataclass(ssd_stats):
            ssd_stats = asdict(ssd_stats)
        elif hasattr(ssd_stats, "to_dict"):
            ssd_stats = ssd_stats.to_dict()
        elif not isinstance(ssd_stats, dict):
            ssd_stats = {}

        ssd_manager = getattr(scheduler, "paged_ssd_cache_manager", None)
        scheduler_model_name = getattr(
            getattr(scheduler, "config", None), "model_name", ""
        )
        if ssd_manager is not None and hasattr(ssd_manager, "get_stats_for_model"):
            try:
                scoped_ssd_stats = ssd_manager.get_stats_for_model(
                    scheduler_model_name or model_id
                )
                if is_dataclass(scoped_ssd_stats):
                    ssd_stats = asdict(scoped_ssd_stats)
                elif isinstance(scoped_ssd_stats, dict):
                    ssd_stats = scoped_ssd_stats
            except Exception as exc:
                logger.warning(
                    "Failed to collect model-scoped SSD cache stats for model '%s': %s",
                    model_id,
                    exc,
                )

        prefix_stats = runtime_stats.get("prefix_cache")
        if is_dataclass(prefix_stats):
            prefix_stats = asdict(prefix_stats)
        elif hasattr(prefix_stats, "to_dict"):
            prefix_stats = prefix_stats.to_dict()
        elif not isinstance(prefix_stats, dict):
            prefix_stats = {}

        indexed_blocks_value = indexed_blocks if isinstance(indexed_blocks, int) else 0
        if not isinstance(block_size, int) or block_size <= 0:
            block_size = int(prefix_stats.get("block_size", 0) or 0)

        partial_block_skips = int(prefix_stats.get("partial_block_skips", 0) or 0)
        partial_tokens_skipped = int(prefix_stats.get("partial_tokens_skipped", 0) or 0)
        last_partial_tokens_skipped = int(
            prefix_stats.get("last_partial_tokens_skipped", 0) or 0
        )
        last_tokens_to_next_block = int(
            prefix_stats.get("last_tokens_to_next_block", 0) or 0
        )

        has_sub_block_cache = (
            indexed_blocks_value == 0
            and isinstance(block_size, int)
            and block_size > 0
            and partial_block_skips > 0
        )

        model_payload = {
            "id": model_id,
            "block_size": block_size,
            "indexed_blocks": indexed_blocks_value,
            "indexed_blocks_display": (
                f"<{block_size}" if has_sub_block_cache else str(indexed_blocks_value)
            ),
            "has_sub_block_cache": has_sub_block_cache,
            "partial_block_skips": partial_block_skips,
            "partial_tokens_skipped": partial_tokens_skipped,
            "last_partial_tokens_skipped": last_partial_tokens_skipped,
            "last_tokens_to_next_block": last_tokens_to_next_block,
            "num_files": int(ssd_stats.get("num_files", 0) or 0),
            "total_size_bytes": int(ssd_stats.get("total_size_bytes", 0) or 0),
            "max_size_bytes": int(ssd_stats.get("max_size_bytes", 0) or 0),
            "hot_cache_max_bytes": int(ssd_stats.get("hot_cache_max_bytes", 0) or 0),
            "hot_cache_size_bytes": int(ssd_stats.get("hot_cache_size_bytes", 0) or 0),
            "hot_cache_entries": int(ssd_stats.get("hot_cache_entries", 0) or 0),
        }

        cache_rates = runtime_stats.get("cache_rates")
        if cache_rates:
            model_payload["cache_rates"] = cache_rates

        payload["models"].append(model_payload)
        payload["total_num_files"] += model_payload["num_files"]
        payload["total_size_bytes"] += model_payload["total_size_bytes"]

        if isinstance(block_size, int) and block_size > 0:
            block_sizes.add(block_size)

    payload["effective_block_sizes"] = sorted(block_sizes)

    # Aggregate hot-cache and disk-max across models. Hot cache max is a single
    # process-wide budget shared by all loaded model managers, so keep the
    # largest reported cap instead of summing per-model rows. Disk max also
    # keeps the config fallback via max() because a single SSD cache directory
    # is shared — the effective cap is the largest configured limit, not a
    # per-model sum.
    hot_cache_max = 0
    disk_max = payload["disk_max_bytes"]
    hot_cache_size_total = 0
    hot_cache_entries_total = 0
    for m in payload["models"]:
        hot_cache_size_total += m.get("hot_cache_size_bytes", 0)
        hot_cache_entries_total += m.get("hot_cache_entries", 0)
        hot_cache_max = max(hot_cache_max, m.get("hot_cache_max_bytes", 0))
        disk_max = max(disk_max, m.get("max_size_bytes", 0))
    payload["hot_cache_max_bytes"] = hot_cache_max
    payload["hot_cache_size_bytes"] = hot_cache_size_total
    payload["hot_cache_entries"] = hot_cache_entries_total
    payload["disk_max_bytes"] = disk_max

    # Fallback: if no loaded models contributed stats, scan the cache
    # directory directly so the dashboard still shows real disk usage.
    if payload["total_num_files"] == 0 and cache_dir.exists():
        try:
            num_files = 0
            total_bytes = 0
            for subdir in "0123456789abcdef":
                subdir_path = cache_dir / subdir
                if not subdir_path.exists():
                    continue
                for f in subdir_path.glob("*.safetensors"):
                    num_files += 1
                    total_bytes += f.stat().st_size
            payload["total_num_files"] = num_files
            payload["total_size_bytes"] = total_bytes
        except Exception as exc:
            logger.warning("Failed to scan SSD cache directory: %s", exc)

    return payload


@router.get("/api/stats")
async def get_server_stats(
    model: str = "",
    scope: str = "session",
    is_admin: bool = Depends(require_admin),
):
    """Get server serving stats for the Status dashboard.

    Args:
        model: Filter by model ID. Empty string returns global aggregate.
        scope: "session" for current session, "alltime" for persisted totals.
    """
    from ..server import resolve_model_id
    from ..server_metrics import get_server_metrics

    metrics = get_server_metrics()
    resolved_model = resolve_model_id(model) or model if model else ""
    snapshot = metrics.get_snapshot(model_id=resolved_model, scope=scope)

    global_settings = _get_global_settings()
    host = global_settings.server.host if global_settings else "127.0.0.1"
    port = global_settings.server.port if global_settings else 8000
    api_key = global_settings.auth.api_key if global_settings else ""

    from ..utils.install import get_cli_prefix

    # Build active_models data for the dashboard card.
    active_models_data = _build_active_models_data()
    runtime_cache_data = _build_runtime_cache_observability(
        global_settings,
        model_filter=model,
    )

    return {
        **snapshot,
        "host": host,
        "port": port,
        "api_key": api_key or "",
        "cli_prefix": get_cli_prefix(),
        "engines": _get_engine_info(),
        "active_models": active_models_data,
        "runtime_cache": runtime_cache_data,
    }


@router.get("/api/activity")
async def get_server_activity(is_admin: bool = Depends(require_admin)):
    """Return lightweight current model and request activity for live displays."""
    return {"active_models": _build_active_models_data()}


def _build_active_models_data() -> dict:
    """Build active models status for the dashboard Active Models card."""
    from ..model_discovery import format_size
    from ..prefill_progress import get_prefill_tracker

    engine_pool = _get_engine_pool()
    server_state = _get_server_state()
    if engine_pool is None:
        return {
            "models": [],
            "model_memory_used": 0,
            "model_memory_max": 0,
            "memory_pressure": {
                "enabled": False,
                "current_bytes": 0,
                "soft_bytes": 0,
                "hard_bytes": 0,
                "current_formatted": "0.0GB",
                "soft_formatted": "0.0GB",
                "hard_formatted": "0.0GB",
                "pressure_level": "ok",
            },
            "total_active_requests": 0,
            "total_waiting_requests": 0,
        }

    now = time.monotonic()
    tracker = get_prefill_tracker()
    status = engine_pool.get_status()
    enforcer = (
        getattr(server_state, "process_memory_enforcer", None)
        if server_state is not None
        else None
    )
    enforcer_status = None
    if enforcer is not None:
        try:
            enforcer_status = enforcer.get_status()
        except Exception as exc:  # noqa: BLE001
            logger.warning("Memory enforcer status unavailable: %s", exc)
    models = []
    total_active = 0
    total_waiting = 0

    for model_info in status.get("models", []):
        if not model_info.get("loaded") and not model_info.get("is_loading"):
            continue

        model_id = model_info["id"]
        active_requests = 0
        waiting_requests = 0
        running_by_id = {}
        has_scheduler_snapshot = False
        waiting_ids = set()
        waiting = []
        activities = []

        # Get per-model active/waiting request counts.
        # Follow the same pattern as server.py /api/status endpoint.
        collector_request_ids: set = set()
        active_request_ids: set = set()
        activity_requests = 0
        entry = engine_pool._entries.get(model_id)
        if entry and entry.engine is not None:
            sched = None
            async_core = getattr(entry.engine, "_engine", None)
            if async_core is not None:
                core = getattr(async_core, "engine", None)
                if core is not None:
                    collectors = getattr(core, "_output_collectors", {})
                    try:
                        collector_request_ids = set(collectors.keys())
                    except RuntimeError:
                        # Scheduler state is mutated from the engine executor;
                        # keep the dashboard endpoint best-effort rather than
                        # failing on a concurrent dict resize.
                        collector_request_ids = set()

                    sched = getattr(core, "scheduler", None)
            else:
                # Engines without an AsyncEngineCore (DFlash) still expose a
                # scheduler once their fallback engine is active.
                sched = getattr(entry.engine, "scheduler", None)
            if sched is not None and hasattr(sched, "snapshot_for_admin"):
                snap = sched.snapshot_for_admin()
                has_scheduler_snapshot = True
                running_by_id = snap["running_by_id"]
                waiting_queue = snap["waiting"]
                waiting_requests = len(waiting_queue)
                waiting_ids = {req.request_id for req in waiting_queue}
                waiting = [
                    {
                        "request_id": req.request_id,
                        "queue_position": idx,
                        "elapsed_seconds": max(0.0, now - req.arrival_time),
                        "prompt_tokens": getattr(req, "num_prompt_tokens", 0),
                    }
                    for idx, req in enumerate(waiting_queue, start=1)
                ]
            if hasattr(entry.engine, "get_activity_snapshot"):
                # Requests the engine tracks itself (non-streaming engines,
                # DFlash primary mode). Counted on top of any scheduler
                # snapshot; the two sources never overlap.
                snapshot = entry.engine.get_activity_snapshot()
                activity_requests = snapshot.get("active_requests", 0)
                activities = snapshot.get("activities", [])

        prefilling = tracker.get_model_progress(model_id)
        prefilling_ids = {p["request_id"] for p in prefilling}
        if has_scheduler_snapshot:
            active_request_ids = set(running_by_id) | prefilling_ids
        elif collector_request_ids:
            active_request_ids = collector_request_ids - waiting_ids
        if has_scheduler_snapshot or collector_request_ids:
            active_requests = len(active_request_ids)
        active_requests += activity_requests

        # Generating = active requests that finished prefill.
        generating = []
        for rid in sorted(active_request_ids - prefilling_ids - waiting_ids):
            req = running_by_id.get(rid)
            generated_tokens = getattr(req, "num_output_tokens", 0) if req else 0
            started_at = getattr(req, "generation_started_at", None) if req else None
            last_activity_at = getattr(req, "last_activity_at", None) if req else None
            elapsed = max(0.0, now - started_at) if started_at else None
            last_activity_age = (
                max(0.0, now - last_activity_at) if last_activity_at else None
            )
            tokens_per_second = (
                generated_tokens / elapsed if elapsed and elapsed > 0 else 0.0
            )
            generating.append(
                {
                    "request_id": rid,
                    "elapsed_seconds": elapsed,
                    "generated_tokens": generated_tokens,
                    "tokens_per_second": tokens_per_second,
                    "last_activity_age_seconds": last_activity_age,
                    "prompt_tokens": getattr(req, "num_prompt_tokens", 0) if req else 0,
                    "max_tokens": getattr(req, "max_tokens", None) if req else None,
                }
            )

        loading_started_at = model_info.get("loading_started_at")
        loading_elapsed_seconds = (
            max(0.0, now - loading_started_at) if loading_started_at else None
        )
        loading_estimated_seconds = None
        loading_remaining_seconds_estimate = None
        if loading_elapsed_seconds is not None:
            estimated_size_gb = model_info.get("estimated_size", 0) / (1024**3)
            # Model loaders do not expose byte-level progress, so use a
            # deliberately conservative elapsed-time estimate and cap below
            # complete until the model is actually loaded.
            observed_seconds_per_gb = status.get("load_seconds_per_gb_estimate")
            observations = status.get("load_time_observations", 0)
            if observed_seconds_per_gb and observations >= 2:
                # Adapt to this machine/session once we have more than a
                # single potentially-misleading sample.
                loading_estimated_seconds = max(
                    3.0,
                    1.0 + estimated_size_gb * float(observed_seconds_per_gb),
                )
                if loading_elapsed_seconds < loading_estimated_seconds:
                    loading_remaining_seconds_estimate = max(
                        0.0, loading_estimated_seconds - loading_elapsed_seconds
                    )

        # Compute idle time and TTL remaining for loaded models.
        is_loaded = (
            model_info.get("loaded") and entry is not None and entry.engine is not None
        )
        last_access = model_info.get("last_access")
        idle_seconds: float | None = None
        ttl_remaining_seconds: float | None = None

        if is_loaded and last_access is not None and last_access > 0:
            idle_seconds = max(0.0, time.time() - last_access)

        # Determine effective TTL: per-model ttl_seconds first, then global idle_timeout.
        effective_ttl: int | None = None
        settings_manager = _get_settings_manager()
        if is_loaded and settings_manager is not None:
            model_settings = settings_manager.get_settings(model_id)
            if (
                model_settings is not None
                and getattr(model_settings, "ttl_seconds", None) is not None
            ):
                effective_ttl = model_settings.ttl_seconds
        if effective_ttl is None:
            global_settings = _get_global_settings()
            if global_settings is not None:
                gt = getattr(global_settings, "idle_timeout", None)
                if gt is not None:
                    effective_ttl = getattr(gt, "idle_timeout_seconds", None)

        if is_loaded and effective_ttl is not None and idle_seconds is not None:
            ttl_remaining_seconds = max(0.0, effective_ttl - idle_seconds)

        # DFlash observability (issue #2398): session speculation counters and
        # the load-time precision pairing warning. None on non-DFlash engines.
        dflash_info = None
        if entry is not None and entry.engine is not None:
            pairing = getattr(entry.engine, "pairing_warning", None)
            speculation = None
            get_speculation = getattr(entry.engine, "get_speculation_stats", None)
            if callable(get_speculation):
                try:
                    speculation = get_speculation()
                except Exception:
                    logger.debug("get_speculation_stats failed", exc_info=True)
            if speculation is not None or pairing:
                dflash_info = {
                    "speculation": speculation,
                    "pairing_warning": pairing,
                }

        models.append(
            {
                "id": model_id,
                "estimated_size": model_info.get("estimated_size", 0),
                "estimated_size_formatted": format_size(
                    model_info.get("estimated_size", 0)
                ),
                "actual_size": model_info.get("actual_size") or 0,
                "actual_size_formatted": (
                    format_size(model_info.get("actual_size", 0))
                    if model_info.get("actual_size")
                    else None
                ),
                "pinned": model_info.get("pinned", False),
                "is_loading": model_info.get("is_loading", False),
                "loading_elapsed_seconds": loading_elapsed_seconds,
                "loading_estimated_seconds": loading_estimated_seconds,
                "loading_remaining_seconds_estimate": loading_remaining_seconds_estimate,
                "active_requests": active_requests,
                "waiting_requests": waiting_requests,
                "waiting": waiting,
                "activities": activities,
                "prefilling": prefilling,
                "generating": generating,
                "idle_seconds": idle_seconds,
                "ttl_remaining_seconds": ttl_remaining_seconds,
                "dflash": dflash_info,
            }
        )

        total_active += active_requests
        total_waiting += waiting_requests

    # model_memory_used reports phys_footprint (whole process) when the
    # enforcer is running so the UI's usage bar matches the value used to
    # drive eviction. model_memory_max is the final_ceiling from
    # enforcer.get_final_ceiling().
    if enforcer_status is not None and enforcer_status.get("enabled"):
        memory_used = enforcer_status.get("current_bytes", 0)
        memory_max = enforcer_status.get("ceiling_bytes", 0)
    else:
        memory_used = status.get("current_model_memory", 0)
        memory_max = status.get("final_ceiling", 0)
    return {
        "models": models,
        "model_memory_used": memory_used,
        "model_memory_max": memory_max,
        "memory_pressure": {
            "enabled": bool(enforcer_status and enforcer_status.get("enabled")),
            "current_bytes": (
                enforcer_status.get("current_bytes", 0)
                if enforcer_status is not None
                else 0
            ),
            "soft_bytes": (
                enforcer_status.get("soft_bytes", 0)
                if enforcer_status is not None
                else 0
            ),
            "hard_bytes": (
                enforcer_status.get("hard_bytes", 0)
                if enforcer_status is not None
                else 0
            ),
            "current_formatted": (
                enforcer_status.get("current_formatted", "0.0GB")
                if enforcer_status is not None
                else "0.0GB"
            ),
            "soft_formatted": (
                enforcer_status.get("soft_formatted", "0.0GB")
                if enforcer_status is not None
                else "0.0GB"
            ),
            "hard_formatted": (
                enforcer_status.get("hard_formatted", "0.0GB")
                if enforcer_status is not None
                else "0.0GB"
            ),
            "pressure_level": (
                enforcer_status.get("pressure_level", "ok")
                if enforcer_status is not None
                else "ok"
            ),
        },
        "total_active_requests": total_active,
        "total_waiting_requests": total_waiting,
    }


@router.post("/api/stats/clear")
async def clear_server_stats(is_admin: bool = Depends(require_admin)):
    """Clear session server metrics."""
    from ..server_metrics import get_server_metrics

    get_server_metrics().clear_metrics()
    return {"status": "ok"}


@router.post("/api/stats/clear-alltime")
async def clear_alltime_stats(is_admin: bool = Depends(require_admin)):
    """Clear all-time server metrics and delete persisted stats file."""
    from ..server_metrics import get_server_metrics

    get_server_metrics().clear_alltime_metrics()
    return {"status": "ok"}


def _iter_loaded_scheduler_records():
    """Yield (model_id, scheduler, core) for each loaded model.

    Traverses the internal engine hierarchy: pool entry → async engine →
    core engine → scheduler.
    """
    engine_pool = _get_engine_pool()
    if engine_pool is None:
        return
    for model_info in engine_pool.get_status().get("models", []):
        model_id = model_info.get("id")
        if not model_id or not model_info.get("loaded"):
            continue
        entry = engine_pool._entries.get(model_id)
        if entry is None or entry.engine is None:
            continue
        async_core = getattr(entry.engine, "_engine", None)
        core = getattr(async_core, "engine", None) if async_core is not None else None
        scheduler = getattr(core, "scheduler", None) if core is not None else None
        if scheduler is not None:
            yield model_id, scheduler, core


def _iter_loaded_schedulers():
    """Yield (model_id, scheduler) for each loaded model.

    Both ``clear_ssd_cache`` and ``clear_hot_cache`` share this traversal.
    """
    for model_id, scheduler, _core in _iter_loaded_scheduler_records():
        yield model_id, scheduler


@router.post("/api/ssd-cache/clear")
async def clear_ssd_cache(is_admin: bool = Depends(require_admin)):
    """Clear all SSD cache files for all loaded models.

    Uses loaded models' SSD cache managers when available.  Falls back to
    direct filesystem deletion so caches can be wiped even when no model
    is loaded.
    """
    total_deleted = 0

    for model_id, scheduler in _iter_loaded_schedulers():
        ssd_manager = getattr(scheduler, "paged_ssd_cache_manager", None)
        if ssd_manager is not None:
            try:
                total_deleted += ssd_manager.clear()
            except Exception as exc:
                logger.warning(
                    "Failed to clear SSD cache for model '%s': %s",
                    model_id,
                    exc,
                )

    # Phase 2: remove any remaining files on disk (covers unloaded models)
    global_settings = _get_global_settings()
    if global_settings is not None:
        cache_dir = global_settings.cache.get_ssd_cache_dir(
            global_settings.base_path,
        )
        if cache_dir.exists():
            try:
                for subdir in "0123456789abcdef":
                    subdir_path = cache_dir / subdir
                    if not subdir_path.exists():
                        continue
                    for f in subdir_path.glob("*.safetensors"):
                        try:
                            f.unlink()
                            total_deleted += 1
                        except OSError:
                            pass
            except Exception as exc:
                logger.warning("Failed to clean SSD cache directory: %s", exc)

    return {"status": "ok", "total_deleted": total_deleted}


@router.post("/api/hot-cache/clear")
async def clear_hot_cache(is_admin: bool = Depends(require_admin)):
    """Clear the in-memory hot cache and release the buffers it held.

    Dropping hot cache entries releases Python references, but MLX may keep
    now-unused buffers in its allocator pool. Reclaim through the scheduler's
    synchronized clear path so active engine streams and async store-cache
    workers keep the same Metal safety barriers used by generation.
    """
    import gc

    from ..engine_core import get_mlx_executor
    from ..scheduler import _sync_and_clear_cache
    from ..utils.proc_memory import get_phys_footprint

    footprint_before = get_phys_footprint()
    total_cleared = 0
    reclaim_targets = []
    for model_id, scheduler, core in _iter_loaded_scheduler_records():
        ssd_manager = getattr(scheduler, "paged_ssd_cache_manager", None)
        if ssd_manager is not None and hasattr(ssd_manager, "clear_hot_cache"):
            try:
                total_cleared += ssd_manager.clear_hot_cache()
            except Exception as exc:
                logger.warning(
                    "Failed to clear hot cache for model '%s': %s",
                    model_id,
                    exc,
                )
        rate_tracker = getattr(scheduler, "_cache_rate_tracker", None)
        if rate_tracker is not None:
            rate_tracker.clear()
        executor = getattr(core, "_mlx_executor", None)
        if executor is not None:
            reclaim_targets.append(
                (model_id, executor, getattr(scheduler, "_stream", None))
            )

    # Also clear managers orphaned by an abnormal teardown: they hold live
    # hot cache but are no longer attached to a loaded scheduler, so the loop
    # above cannot reach them. The shared budget still references them.
    pool = _get_engine_pool()
    budget = getattr(getattr(pool, "_scheduler_config", None), "hot_cache_budget", None)
    if budget is not None and hasattr(budget, "clear_all_owners"):
        try:
            total_cleared += budget.clear_all_owners()
        except Exception as exc:
            logger.warning("Failed to clear orphaned hot caches: %s", exc)

    # Return pooled buffers to the OS using scheduler._sync_and_clear_cache(),
    # the same lock/synchronize/clear helper used by generation. Run on each
    # loaded engine's executor so its thread-local stream is present. If every
    # model has been unloaded, still run one reclaim on the fallback executor so
    # orphaned/no-loaded hot cache cleanup can release MLX's allocator pool.
    gc.collect()
    loop = asyncio.get_running_loop()
    if reclaim_targets:
        for model_id, executor, stream in reclaim_targets:
            try:
                await loop.run_in_executor(executor, _sync_and_clear_cache, stream)
            except RuntimeError as exc:
                if "cannot schedule new futures after shutdown" not in str(exc):
                    raise
                logger.warning(
                    "Engine executor unavailable while reclaiming MLX buffers "
                    "for model '%s': %s",
                    model_id,
                    exc,
                )
                await loop.run_in_executor(get_mlx_executor(), _sync_and_clear_cache)
    else:
        await loop.run_in_executor(get_mlx_executor(), _sync_and_clear_cache)
    bytes_reclaimed = max(0, footprint_before - get_phys_footprint())

    return {
        "status": "ok",
        "total_cleared": total_cleared,
        "bytes_reclaimed": bytes_reclaimed,
    }


def _normalize_probe_tool_calls(messages: list[dict]) -> list[dict]:
    """Parse echoed tool_call arguments (JSON string -> object) for templating.

    Native tool-calling chat templates (GLM, Qwen3.x, MiniMax) iterate
    ``tool_call.function.arguments.items()``, but the OpenAI wire form sends
    ``arguments`` as a JSON string. Rendering the string form raises
    ``'str object' has no attribute 'items'`` and the probe 400s, so any
    conversation that used tools reports an error (hollow cache dot) instead
    of a real hit/miss. The chat path parses these before rendering; mirror
    that here so (a) tool conversations tokenize and (b) the probe's block
    hashes line up with what a real prefill produced. Returns shallow copies
    so the caller's message dicts are left untouched.
    """
    normalized: list[dict] = []
    for msg in messages:
        tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
        if not tool_calls:
            normalized.append(msg)
            continue
        new_calls = []
        for tc in tool_calls:
            fn = tc.get("function") if isinstance(tc, dict) else None
            if isinstance(fn, dict) and "arguments" in fn:
                arguments = _coerce_tool_call_arguments(fn["arguments"])
                tc = {
                    **tc,
                    "function": {**fn, "arguments": _try_parse_json(arguments)},
                }
            new_calls.append(tc)
        normalized.append({**msg, "tool_calls": new_calls})
    return normalized


def _probe_chat_template_kwargs(
    request: "CacheProbeRequest",
    *,
    preserve_thinking_default: bool | None = None,
) -> dict | None:
    """Chat-template kwargs the scheduler would actually prefill this with.

    The probe answers "is this prompt cached", so it has to render byte-for
    byte what a real turn renders. Rendering with the caller's kwargs alone
    ignores the model's own settings — a model with enable_thinking set (or
    any forced/persisted chat_template_kwargs) then hashes a prompt that is
    never prefilled, and since the block walk stops at the first miss, every
    block reports cold.
    """
    settings = None
    if _get_settings_manager is not None:
        try:
            manager = _get_settings_manager()
            if manager is not None:
                settings = manager.get_settings_for_request(
                    request.model_id,
                    resolved_model_id=request.model_id,
                )
        except Exception:
            # A settings lookup failure must not break probing outright —
            # fall back to the caller's kwargs (pre-fix behaviour).
            logger.warning(
                "cache probe: model settings lookup failed for %s; "
                "rendering with request kwargs only",
                request.model_id,
                exc_info=True,
            )
            settings = None
    return (
        merge_chat_template_kwargs(
            settings,
            request.chat_template_kwargs,
            thinking_budget=request.thinking_budget,
            preserve_thinking_default=preserve_thinking_default,
        )
        or None
    )


@router.post("/api/cache/probe")
async def probe_cache(
    request: CacheProbeRequest,
    is_admin: bool = Depends(require_admin),
):
    """Probe cache state for a chat message list.

    Classifies each block of the rendered prompt into one of three buckets:
    - ``blocks_ssd_hot``: in the SSD manager's hot cache (RAM copy of cold
      blocks, ready to mount without disk read)
    - ``blocks_ssd_disk``: only in the SSD index on disk
    - ``blocks_cold``: not cached anywhere (requires full prefill)

    The split is computed via a walk of the chain-hashed block sequence — the
    same hashing the scheduler uses at prefill time. The model must be loaded
    for the probe to run; unloaded models return ``model_loaded: false``.
    """
    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    entry = engine_pool._entries.get(request.model_id)
    if entry is None:
        raise HTTPException(
            status_code=404, detail=f"Model not found: {request.model_id}"
        )
    if entry.engine is None:
        return {
            "model_id": request.model_id,
            "model_loaded": False,
            "reason": "Model is not loaded — load it to enable cache probing.",
        }

    engine = entry.engine
    tokenizer = getattr(engine, "_tokenizer", None)
    if tokenizer is None or not hasattr(tokenizer, "apply_chat_template"):
        raise HTTPException(
            status_code=400,
            detail="Model tokenizer does not support chat templating.",
        )

    # Reach into the scheduler to access the prefix index and SSD manager.
    async_core = getattr(engine, "_engine", None)
    core = getattr(async_core, "engine", None) if async_core is not None else None
    scheduler = getattr(core, "scheduler", None) if core is not None else None
    if scheduler is None:
        raise HTTPException(
            status_code=500, detail="Scheduler unavailable for loaded model."
        )

    prefix_cache = getattr(scheduler, "block_aware_cache", None)
    ssd_manager = getattr(scheduler, "paged_ssd_cache_manager", None)
    paged_cache = getattr(scheduler, "paged_cache_manager", None)
    block_size = getattr(
        getattr(scheduler, "config", None), "paged_cache_block_size", 0
    )
    if not block_size and prefix_cache is not None:
        block_size = getattr(prefix_cache, "block_size", 0)
    if not block_size:
        raise HTTPException(
            status_code=500,
            detail="Cache block size unavailable — cache may not be enabled.",
        )

    # Render + tokenize the prompt using the same path as generation so the
    # hashes line up with what the scheduler would produce at prefill.
    try:
        messages = _normalize_probe_tool_calls(request.messages)
        if hasattr(engine, "_preprocess_messages"):
            messages = engine._preprocess_messages(messages)
        try:
            from ..api.tool_calling import convert_tools_for_template  # type: ignore

            template_tools = (
                convert_tools_for_template(request.tools) if request.tools else None
            )
        except Exception:
            template_tools = request.tools or None
        if hasattr(engine, "_apply_chat_template"):
            prompt = engine._apply_chat_template(
                messages,
                template_tools,
                chat_template_kwargs=_probe_chat_template_kwargs(
                    request,
                    preserve_thinking_default=getattr(
                        entry, "preserve_thinking_default", None
                    ),
                ),
            )
        else:
            prompt = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
        token_ids = list(tokenizer.encode(prompt))
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail=f"Failed to tokenize messages: {exc}"
        )

    total_tokens = len(token_ids)
    if total_tokens == 0:
        return {
            "model_id": request.model_id,
            "model_loaded": True,
            "total_tokens": 0,
            "block_size": block_size,
            "total_blocks": 0,
            "blocks_ssd_hot": 0,
            "blocks_ssd_disk": 0,
            "blocks_cold": 0,
            "ssd_hit_tokens": 0,
            "cold_tokens": 0,
        }

    # Compute chain-hashed block sequence.
    from ..cache.paged_cache import compute_block_hash

    model_name = getattr(paged_cache, "model_name", None) if paged_cache else None
    ssd_index = getattr(ssd_manager, "_index", None) if ssd_manager else None
    ssd_hot = getattr(ssd_manager, "_hot_cache", None) if ssd_manager else None

    # The cache is a contiguous prefix (each block chain-hashed from the
    # previous), so we walk block-by-block until the first retrievability
    # miss — after that, every subsequent block is necessarily cold.
    #
    # Ground truth for "cached" in paged-SSD mode is retrievability:
    # hot_cache (RAM copy) OR ssd_index (on disk). BlockAwarePrefixCache's
    # internal prefix index is deliberately NOT consulted — it tracks every
    # hash the scheduler has seen and isn't cleared by clear_ssd_cache(),
    # so relying on it would report false positives after a manual wipe.
    blocks_ssd_hot = 0
    blocks_ssd_disk = 0
    ssd_hit_tokens = 0

    parent_hash = b""
    total_blocks = (total_tokens + block_size - 1) // block_size

    for start in range(0, total_tokens, block_size):
        end = min(start + block_size, total_tokens)
        block_tokens = token_ids[start:end]
        if not block_tokens:
            break

        block_hash = compute_block_hash(
            parent_hash,
            block_tokens,
            extra_keys=None,
            model_name=model_name,
        )
        parent_hash = block_hash

        in_ssd_hot = ssd_hot is not None and block_hash in ssd_hot
        in_ssd_disk = False
        if ssd_index is not None:
            try:
                in_ssd_disk = ssd_index.contains(block_hash)
            except Exception:
                in_ssd_disk = False

        if not (in_ssd_hot or in_ssd_disk):
            break

        if in_ssd_hot:
            blocks_ssd_hot += 1
        else:
            blocks_ssd_disk += 1
        ssd_hit_tokens += len(block_tokens)

    cached_blocks = blocks_ssd_hot + blocks_ssd_disk
    blocks_cold = max(total_blocks - cached_blocks, 0)

    return {
        "model_id": request.model_id,
        "model_loaded": True,
        "total_tokens": total_tokens,
        "block_size": block_size,
        "total_blocks": total_blocks,
        "blocks_ssd_hot": blocks_ssd_hot,
        "blocks_ssd_disk": blocks_ssd_disk,
        "blocks_cold": blocks_cold,
        "ssd_hit_tokens": ssd_hit_tokens,
        "cold_tokens": max(total_tokens - ssd_hit_tokens, 0),
    }


# =============================================================================
# HuggingFace Downloader API Routes
# =============================================================================


@router.post("/api/hf/download")
async def start_hf_download(
    request: HFDownloadRequest,
    is_admin: bool = Depends(require_admin),
):
    """Start downloading a model from HuggingFace."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    try:
        task = await _hf_downloader.start_download(request.repo_id, request.hf_token)
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.get("/api/hf/tasks")
async def list_hf_tasks(is_admin: bool = Depends(require_admin)):
    """List all download tasks."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    return {"tasks": _hf_downloader.get_tasks()}


@router.post("/api/hf/cancel/{task_id}")
async def cancel_hf_download(
    task_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Cancel an active download."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    success = await _hf_downloader.cancel_download(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or not cancellable")
    return {"success": True}


@router.post("/api/hf/retry/{task_id}")
async def retry_hf_download(
    task_id: str,
    request: HFRetryRequest = HFRetryRequest(),
    is_admin: bool = Depends(require_admin),
):
    """Retry a failed or cancelled download, resuming from existing files."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    try:
        task = await _hf_downloader.retry_download(task_id, request.hf_token)
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.delete("/api/hf/task/{task_id}")
async def remove_hf_task(
    task_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Remove a completed, failed, or cancelled task."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    success = _hf_downloader.remove_task(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or still active")
    return {"success": True}


@router.get("/api/hf/recommended")
async def get_recommended_models(
    mlx_only: bool = True,
    is_admin: bool = Depends(require_admin),
):
    """Get recommended models filtered by system memory."""
    if _hf_downloader is None:
        raise HTTPException(status_code=503, detail="Downloader not initialized")

    memory_info = get_system_memory_info()
    max_memory = memory_info["total_bytes"] or 16 * 1024**3

    from .hf_downloader import HFDownloader

    try:
        result = await HFDownloader.get_recommended_models(
            max_memory_bytes=max_memory, result_limit=50, mlx_only=mlx_only
        )
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="HuggingFace API request timed out. The service may be temporarily unavailable.",
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))


@router.get("/api/hf/search")
async def search_hf_models(
    q: str = "",
    sort: str = "trending",
    limit: int = 100,
    mlx_only: bool = True,
    # Filtering
    min_params: Optional[int] = None,
    max_params: Optional[int] = None,
    min_size: Optional[int] = None,  # bytes
    max_size: Optional[int] = None,  # bytes
    # Sorting
    sort_by_size: bool = False,
    sort_ascending: bool = False,
    is_admin: bool = Depends(require_admin),
):
    """Search HuggingFace models by query with filtering and sorting.

    Query Parameters:
        q: Search query string (required)
        sort: Sort order - trending/downloads/created/updated/most_params/least_params/largest/smallest
        limit: Maximum results (max 100)
        mlx_only: Restrict to MLX library models
        min_params: Minimum parameter count
        max_params: Maximum parameter count
        min_size: Minimum model size in bytes
        max_size: Maximum model size in bytes
        sort_by_size: Sort results by size instead of default sort
        sort_ascending: Sort in ascending order
    """
    if not q.strip():
        raise HTTPException(status_code=400, detail="Query parameter 'q' is required")

    from .hf_downloader import HFDownloader

    try:
        result = await HFDownloader.search_models(
            query=q.strip(),
            sort=sort,
            limit=min(limit, 100),
            mlx_only=mlx_only,
            min_params=min_params,
            max_params=max_params,
            min_size=min_size,
            max_size=max_size,
            sort_by_size=sort_by_size,
            sort_ascending=sort_ascending,
        )
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="HuggingFace API request timed out. The service may be temporarily unavailable.",
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))


@router.get("/api/hf/model-info")
async def get_hf_model_info(
    repo_id: str = "",
    is_admin: bool = Depends(require_admin),
):
    """Get detailed model information from HuggingFace."""
    if not repo_id.strip():
        raise HTTPException(
            status_code=400, detail="Query parameter 'repo_id' is required"
        )

    from huggingface_hub.utils import RepositoryNotFoundError

    from .hf_downloader import HFDownloader

    try:
        result = await HFDownloader.get_model_info(repo_id=repo_id.strip())
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="HuggingFace API request timed out. The service may be temporarily unavailable.",
        )
    except RepositoryNotFoundError:
        raise HTTPException(
            status_code=404, detail=f"Model '{repo_id.strip()}' not found"
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))


@router.get("/api/hf/models")
async def list_hf_models(is_admin: bool = Depends(require_admin)):
    """List models in all model directories with disk size info."""
    global_settings = _get_global_settings()
    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    model_dirs = global_settings.model.get_model_dirs(global_settings.base_path)

    from ..model_discovery import _resolve_hf_cache_entry

    def _add_model(
        model_path: Path,
        model_name: str,
        *,
        source_repo_id: str | None = None,
    ) -> None:
        if model_name in seen_names:
            return
        seen_names.add(model_name)
        total_size = sum(f.stat().st_size for f in model_path.rglob("*") if f.is_file())
        models.append(
            {
                "name": model_name,
                "path": str(model_path),
                "display_name": _model_display_name(
                    model_name,
                    model_path,
                    model_dirs,
                    source_repo_id=source_repo_id,
                ),
                "size": total_size,
                "size_formatted": format_size(total_size),
            }
        )

    models = []
    seen_names: set[str] = set()
    for model_dir in model_dirs:
        if not model_dir.exists():
            continue
        for subdir in sorted(model_dir.iterdir()):
            if not subdir.is_dir() or subdir.name.startswith("."):
                continue

            if (subdir / "config.json").exists():
                # Level 1: direct model folder
                _add_model(subdir, subdir.name)
            else:
                # HF Hub cache entry: models--Org--Name/snapshots/<hash>/
                hf_resolved = _resolve_hf_cache_entry(subdir)
                if hf_resolved is not None:
                    if (hf_resolved.snapshot_path / "config.json").exists():
                        _add_model(
                            hf_resolved.snapshot_path,
                            hf_resolved.model_id,
                            source_repo_id=hf_resolved.source_repo_id,
                        )
                    continue

                # Level 2: organization folder — scan children
                for child in sorted(subdir.iterdir()):
                    if not child.is_dir() or child.name.startswith("."):
                        continue
                    if (child / "config.json").exists():
                        _add_model(child, child.name)

    # Sort by the UI display name so organization prefixes group together.
    models.sort(key=lambda m: m["display_name"].lower())
    return {"models": models}


@router.delete("/api/hf/models/{model_name}")
async def delete_hf_model(
    model_name: str,
    is_admin: bool = Depends(require_admin),
):
    """Delete a downloaded model from disk and refresh the model pool."""
    global_settings = _get_global_settings()
    engine_pool = _get_engine_pool()

    if global_settings is None:
        raise HTTPException(status_code=503, detail="Server not initialized")

    model_dirs = global_settings.model.get_model_dirs(global_settings.base_path)

    # Search for model across all directories in both flat and org-folder layouts
    model_path = None
    parent_model_dir = None
    for model_dir in model_dirs:
        if not model_dir.exists():
            continue
        candidate = model_dir / model_name
        if candidate.is_dir() and (candidate / "config.json").exists():
            model_path = candidate
            parent_model_dir = model_dir
            break
        # Try two-level: search inside organization folders
        for subdir in model_dir.iterdir():
            if not subdir.is_dir() or subdir.name.startswith("."):
                continue
            candidate = subdir / model_name
            if candidate.is_dir() and (candidate / "config.json").exists():
                model_path = candidate
                parent_model_dir = model_dir
                break
        if model_path is not None:
            break

    if model_path is None:
        raise HTTPException(status_code=404, detail="Model not found")

    # Validate path traversal against parent model directory
    try:
        if not model_path.resolve().is_relative_to(parent_model_dir.resolve()):
            raise HTTPException(status_code=400, detail="Invalid model name")
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid model name")

    if not model_path.is_dir():
        raise HTTPException(status_code=400, detail="Not a model directory")

    # Unload model if loaded
    if engine_pool is not None:
        loaded_ids = engine_pool.get_loaded_model_ids()
        if model_name in loaded_ids:
            try:
                await engine_pool._unload_engine(model_name)
                logger.info(f"Unloaded model '{model_name}' before deletion")
            except Exception as e:
                logger.warning(f"Failed to unload model '{model_name}': {e}")

    # Delete from disk
    # Handle macOS resource fork files (._*) that may disappear on non-native
    # filesystems (exFAT, NTFS). Use onexc (Python 3.12+) to avoid
    # DeprecationWarning, with onerror fallback for older versions.
    def _handle_onexc(func, path, exc):
        if isinstance(exc, FileNotFoundError) and Path(path).name.startswith("._"):
            logger.debug(f"Ignoring missing resource fork file: {path}")
            return
        raise exc

    def _handle_onerror(func, path, exc_info):
        if exc_info[0] == FileNotFoundError and Path(path).name.startswith("._"):
            logger.debug(f"Ignoring missing resource fork file: {path}")
            return
        raise exc_info[1].with_traceback(exc_info[2])

    try:
        if sys.version_info >= (3, 12):
            shutil.rmtree(model_path, onexc=_handle_onexc)
        else:
            shutil.rmtree(model_path, onerror=_handle_onerror)
        logger.info(f"Deleted model directory: {model_path}")
    except Exception as e:
        logger.error(f"Failed to delete model directory {model_path}: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to delete model: {e}")

    # If the model was inside an org folder (organized layout) and that
    # folder is now empty, drop it so the listing stays tidy.
    parent = model_path.parent
    if parent != parent_model_dir and parent.exists() and not any(parent.iterdir()):
        try:
            parent.rmdir()
            logger.info(f"Removed empty org folder: {parent}")
        except OSError as e:
            logger.debug(f"Could not remove empty org folder {parent}: {e}")

    # Re-discover models
    if engine_pool is not None:
        settings_manager = _get_settings_manager()
        pinned_models = []
        if settings_manager:
            pinned_models = settings_manager.get_pinned_model_ids()

        engine_pool._entries.pop(model_name, None)
        # Release the deleted model's persisted settings (including its alias)
        # so they can be reused by another model.
        if settings_manager:
            settings_manager.delete_settings(model_name)
        engine_pool.discover_models(
            [str(d) for d in global_settings.get_effective_model_dirs()],
            pinned_models,
        )
        if settings_manager:
            engine_pool.apply_settings_overrides(settings_manager)
        logger.info("Model pool refreshed after deletion")

    return {"success": True, "message": f"Model '{model_name}' deleted"}


# =============================================================================
# ModelScope Downloader API Routes
# =============================================================================


@router.get("/api/ms/status")
async def ms_status(is_admin: bool = Depends(require_admin)):
    """Check if ModelScope downloader is available."""
    return {"available": _ms_downloader is not None}


@router.post("/api/ms/download")
async def start_ms_download(
    request: MSDownloadRequest,
    is_admin: bool = Depends(require_admin),
):
    """Start downloading a model from ModelScope."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    try:
        task = await _ms_downloader.start_download(request.model_id, request.ms_token)
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))


@router.get("/api/ms/tasks")
async def list_ms_tasks(is_admin: bool = Depends(require_admin)):
    """List all ModelScope download tasks."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    return {"tasks": _ms_downloader.get_tasks()}


@router.post("/api/ms/cancel/{task_id}")
async def cancel_ms_download(
    task_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Cancel an active ModelScope download."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    success = await _ms_downloader.cancel_download(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or not cancellable")
    return {"success": True}


@router.post("/api/ms/retry/{task_id}")
async def retry_ms_download(
    task_id: str,
    request: MSRetryRequest = MSRetryRequest(),
    is_admin: bool = Depends(require_admin),
):
    """Retry a failed or cancelled ModelScope download."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    try:
        task = await _ms_downloader.retry_download(task_id, request.ms_token)
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.delete("/api/ms/task/{task_id}")
async def remove_ms_task(
    task_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Remove a completed, failed, or cancelled ModelScope task."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    success = _ms_downloader.remove_task(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or still active")
    return {"success": True}


@router.get("/api/ms/recommended")
async def get_ms_recommended_models(
    mlx_only: bool = True,
    is_admin: bool = Depends(require_admin),
):
    """Get recommended models from ModelScope filtered by system memory."""
    if _ms_downloader is None:
        raise HTTPException(
            status_code=503, detail="ModelScope downloader not initialized"
        )

    memory_info = get_system_memory_info()
    max_memory = memory_info["total_bytes"] or 16 * 1024**3

    from .ms_downloader import MSDownloader

    try:
        result = await MSDownloader.get_recommended_models(
            max_memory_bytes=max_memory, result_limit=50, mlx_only=mlx_only
        )
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="ModelScope API request timed out. The service may be temporarily unavailable.",
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))


@router.get("/api/ms/search")
async def search_ms_models(
    q: str = "",
    sort: str = "trending",
    limit: int = 100,
    mlx_only: bool = True,
    is_admin: bool = Depends(require_admin),
):
    """Search ModelScope models by query."""
    if not q.strip():
        raise HTTPException(status_code=400, detail="Query parameter 'q' is required")

    from .ms_downloader import MSDownloader

    try:
        result = await MSDownloader.search_models(
            query=q.strip(),
            sort=sort,
            limit=min(limit, 100),
            mlx_only=mlx_only,
        )
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="ModelScope API request timed out. The service may be temporarily unavailable.",
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e))


@router.get("/api/ms/model-info")
async def get_ms_model_info(
    model_id: str = "",
    is_admin: bool = Depends(require_admin),
):
    """Get detailed model information from ModelScope."""
    if not model_id.strip():
        raise HTTPException(
            status_code=400, detail="Query parameter 'model_id' is required"
        )

    from .ms_downloader import MSDownloader

    try:
        result = await MSDownloader.get_model_info(model_id=model_id.strip())
        return result
    except TimeoutError:
        raise HTTPException(
            status_code=504,
            detail="ModelScope API request timed out. The service may be temporarily unavailable.",
        )
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        if "NotExistError" in type(e).__name__ or "404" in str(e):
            raise HTTPException(
                status_code=404, detail=f"Model '{model_id.strip()}' not found"
            )
        raise HTTPException(status_code=502, detail=str(e))


# =============================================================================
# Accuracy Benchmark API Routes (MUST be before throughput {bench_id} routes)
# =============================================================================


@router.post("/api/bench/accuracy/queue/add")
async def add_to_accuracy_queue(
    request: Request,
    is_admin: bool = Depends(require_admin),
):
    """Add a model to the accuracy benchmark queue and start if idle."""
    from .accuracy_benchmark import (
        AccuracyBenchmarkRequest,
        add_to_queue,
        get_queue_status,
        start_next_from_queue,
    )

    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    from .context_benchmark import get_active_run as get_active_context_run

    context_active = get_active_context_run()
    if context_active is not None:
        raise HTTPException(
            status_code=409,
            detail=(
                f"A context benchmark is already running "
                f"(bench_id={context_active.bench_id}, "
                f"model_id={context_active.request.model_id})."
            ),
        )

    body = await request.json()
    try:
        bench_request = AccuracyBenchmarkRequest(**body)
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

    # External runs target a remote model — nothing to validate locally.
    if bench_request.external is None:
        entry = engine_pool.get_entry(bench_request.model_id)
        if entry is None:
            raise HTTPException(
                status_code=404, detail=f"Model not found: {bench_request.model_id}"
            )
        if entry.model_type not in ("llm", "vlm", None):
            raise HTTPException(
                status_code=400,
                detail=f"Model {bench_request.model_id} is not a supported model (type: {entry.model_type})",
            )

    add_to_queue(bench_request)

    logger.info(
        f"Accuracy queue: added {bench_request.model_id} "
        f"benchmarks={list(bench_request.benchmarks.keys())}"
    )

    # Start processing if not already running (synchronous — sets bench_id immediately)
    start_next_from_queue(engine_pool)

    return get_queue_status()


@router.get("/api/bench/accuracy/queue/status")
async def get_accuracy_queue_status(
    is_admin: bool = Depends(require_admin),
):
    """Get accuracy benchmark queue status."""
    from .accuracy_benchmark import get_queue_status

    return get_queue_status()


@router.delete("/api/bench/accuracy/queue/{idx}")
async def remove_from_accuracy_queue(
    idx: int,
    is_admin: bool = Depends(require_admin),
):
    """Remove an item from the accuracy benchmark queue."""
    from .accuracy_benchmark import get_queue_status, remove_from_queue

    if not remove_from_queue(idx):
        raise HTTPException(status_code=404, detail=f"Queue index {idx} not found")

    return get_queue_status()


@router.get("/api/bench/accuracy/results")
async def get_accumulated_accuracy_results(
    is_admin: bool = Depends(require_admin),
):
    """Get all accumulated accuracy benchmark results."""
    from .accuracy_benchmark import get_accumulated_results, get_queue_status

    status = get_queue_status()
    return {
        "results": get_accumulated_results(),
        "running": status["running"],
        "current_model": status["current_model"],
        "current_bench_id": status["current_bench_id"],
    }


@router.post("/api/bench/accuracy/results/reset")
async def reset_accuracy_results(
    is_admin: bool = Depends(require_admin),
):
    """Clear all accumulated accuracy benchmark results."""
    from .accuracy_benchmark import reset_accumulated_results

    reset_accumulated_results()
    return {"status": "reset"}


@router.post("/api/bench/accuracy/cancel")
async def cancel_accuracy_queue(
    is_admin: bool = Depends(require_admin),
):
    """Cancel the current run and clear the queue."""
    from .accuracy_benchmark import cancel_queue

    await cancel_queue()
    return {"status": "cancelled"}


@router.get("/api/bench/accuracy/{bench_id}/stream")
async def stream_accuracy_benchmark(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Stream accuracy benchmark progress via Server-Sent Events."""
    import json

    from fastapi.responses import StreamingResponse

    from .accuracy_benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(
            status_code=404, detail=f"Accuracy benchmark not found: {bench_id}"
        )

    async def event_generator():
        # Replay-then-attach: every subscriber starts at offset 0 of the
        # run's event log and follows along live. Lets the HTML dashboard
        # recover its view on page refresh and lets multiple consumers
        # (e.g. browser + Swift app) share the same run.
        seen = 0
        try:
            while True:
                async with run.cond:
                    while seen >= len(run.events) and not run.terminal:
                        try:
                            await asyncio.wait_for(run.cond.wait(), timeout=60.0)
                        except TimeoutError:
                            break
                    new = list(run.events[seen:])
                    seen = len(run.events)
                    done = run.terminal

                for ev in new:
                    yield f"data: {json.dumps(ev)}\n\n"
                if not new and not done:
                    yield ": keepalive\n\n"
                if done:
                    break
        except asyncio.CancelledError:
            pass

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


# =============================================================================
# Context Benchmark API Routes (MUST be before throughput {bench_id} routes)
# =============================================================================


@router.get("/api/bench/context/active")
async def get_active_context_benchmark(is_admin: bool = Depends(require_admin)):
    """Return the currently-running context benchmark, if any."""
    from .context_benchmark import get_active_run

    run = get_active_run()
    if run is None:
        return {"running": False, "bench_id": None, "model_id": None}
    return {
        "running": True,
        "bench_id": run.bench_id,
        "model_id": run.request.model_id,
        "target_tokens": run.request.target_tokens,
    }


@router.post("/api/bench/context/start")
async def start_context_benchmark(
    request: Request,
    is_admin: bool = Depends(require_admin),
):
    """Start a context window benchmark run.

    Rejects with 409 while any benchmark (context, throughput, accuracy)
    is running — they all unload/load models and would corrupt each
    other. Rejects with 400 when the memory guard is off: there is no
    admission boundary to measure and an unguarded probe prefill can
    genuinely exhaust the machine.
    """
    from .accuracy_benchmark import get_queue_status
    from .benchmark import get_active_run as get_active_throughput_run
    from .context_benchmark import (
        ContextBenchmarkRequest,
        cleanup_old_runs,
        create_run,
        get_active_run,
        run_context_benchmark,
    )

    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    active = get_active_run()
    if active is not None:
        raise HTTPException(
            status_code=409,
            detail=(
                f"A context benchmark is already running "
                f"(bench_id={active.bench_id}, "
                f"model_id={active.request.model_id})."
            ),
        )
    throughput_active = get_active_throughput_run()
    if throughput_active is not None:
        raise HTTPException(
            status_code=409,
            detail=(
                f"A throughput benchmark is already running "
                f"(bench_id={throughput_active.bench_id}, "
                f"model_id={throughput_active.request.model_id})."
            ),
        )
    accuracy_status = get_queue_status()
    if accuracy_status.get("running"):
        raise HTTPException(
            status_code=409,
            detail=(
                f"An accuracy benchmark is already running "
                f"(model_id={accuracy_status.get('current_model')})."
            ),
        )

    from ..server import _server_state

    enforcer = getattr(_server_state, "process_memory_enforcer", None)
    final_ceiling = 0
    if enforcer is not None:
        try:
            final_ceiling = int(enforcer.get_final_ceiling())
        except Exception:
            final_ceiling = 0
    if final_ceiling <= 0:
        raise HTTPException(
            status_code=400,
            detail=(
                "Memory Guard is disabled. The context benchmark measures "
                "the guard's admission boundary, and probing without it can "
                "exhaust system memory. Enable Memory Guard and retry."
            ),
        )

    body = await request.json()
    try:
        bench_request = ContextBenchmarkRequest(**body)
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

    entry = engine_pool.get_entry(bench_request.model_id)
    if entry is None:
        raise HTTPException(
            status_code=404, detail=f"Model not found: {bench_request.model_id}"
        )
    if entry.model_type not in ("llm", "vlm", None):
        raise HTTPException(
            status_code=400,
            detail=(
                f"Model {bench_request.model_id} is not a supported model "
                f"(type: {entry.model_type})"
            ),
        )

    cleanup_old_runs()
    run = create_run(bench_request)
    run.task = asyncio.create_task(run_context_benchmark(run, engine_pool))

    logger.info(
        f"Context benchmark started: {run.bench_id} "
        f"model={bench_request.model_id} target={bench_request.target_tokens}"
    )

    return {
        "bench_id": run.bench_id,
        "status": "started",
        "target_tokens": bench_request.target_tokens,
    }


@router.get("/api/bench/context/{bench_id}/stream")
async def stream_context_benchmark(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Stream context benchmark progress via Server-Sent Events."""
    import json

    from fastapi.responses import StreamingResponse

    from .context_benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(
            status_code=404, detail=f"Benchmark not found: {bench_id}"
        )

    async def event_generator():
        # Replay-then-attach, same shape as the throughput bench stream.
        # Terminal events here are `done` and `error`.
        seen = 0
        try:
            while True:
                async with run.cond:
                    while seen >= len(run.events) and not run.terminal:
                        try:
                            await asyncio.wait_for(run.cond.wait(), timeout=60.0)
                        except TimeoutError:
                            break
                    new = list(run.events[seen:])
                    seen = len(run.events)
                    done = run.terminal

                for ev in new:
                    yield f"data: {json.dumps(ev)}\n\n"
                if not new and not done:
                    yield ": keepalive\n\n"
                if done:
                    break
        except asyncio.CancelledError:
            pass

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@router.post("/api/bench/context/{bench_id}/cancel")
async def cancel_context_benchmark(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Cancel a running context benchmark."""
    from .context_benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(
            status_code=404, detail=f"Benchmark not found: {bench_id}"
        )

    if run.status != "running":
        raise HTTPException(
            status_code=400,
            detail=f"Benchmark is not running (status: {run.status})",
        )

    if run.task and not run.task.done():
        run.task.cancel()

    return {"status": "cancelled", "bench_id": bench_id}


@router.get("/api/bench/context/{bench_id}/results")
async def get_context_benchmark_results(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Get status and result of a context benchmark (REST poll surface)."""
    from .context_benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(
            status_code=404, detail=f"Benchmark not found: {bench_id}"
        )

    return {
        "bench_id": run.bench_id,
        "status": run.status,
        "phase": run.phase,
        "progress": run.progress,
        "message": run.message,
        "result": run.result,
        "error": run.error_message if run.error_message else None,
    }


# =============================================================================
# Benchmark API Routes (Throughput)
# =============================================================================


@router.get("/api/bench/active")
async def get_active_benchmark(is_admin: bool = Depends(require_admin)):
    """Return the currently-running throughput benchmark, if any.

    Symmetric to `/api/bench/accuracy/queue/status` — lets a fresh page
    load or a second tab discover an in-flight run so it can attach to
    the SSE stream. Combined with the replay-on-subscribe stream this
    is what makes the multi-tab + page-refresh story actually work.
    """
    from .benchmark import get_active_run

    run = get_active_run()
    if run is None:
        return {
            "running": False,
            "bench_id": None,
            "model_id": None,
            "context_profile": None,
        }
    return {
        "running": True,
        "bench_id": run.bench_id,
        "model_id": run.request.model_id,
        "context_profile": run.request.context_profile.value,
        "force_lm_engine": run.request.force_lm_engine,
        # Reconnecting tabs need this to restore the disabled-dropdown UI
        # state. Never expose base_url/api_key here — model_id already
        # carries the external model name.
        "external": run.request.external is not None,
    }


@router.post("/api/bench/start")
async def start_benchmark(
    request: Request,
    is_admin: bool = Depends(require_admin),
):
    """Start a benchmark run.

    Validates the model, creates a benchmark run, and starts it
    as an asyncio background task. Rejects with 409 if another
    throughput bench is already running — two concurrent runs on
    the same engine produce mutually-corrupted measurements.
    """
    from .benchmark import (
        BenchmarkRequest,
        cleanup_old_runs,
        create_run,
        get_active_run,
        run_benchmark,
    )

    engine_pool = _get_engine_pool()
    if engine_pool is None:
        raise HTTPException(status_code=503, detail="Engine pool not initialized")

    # One throughput bench at a time. The replay-on-subscribe stream lets
    # clients attach to the already-running one if that's what they want.
    active = get_active_run()
    if active is not None:
        raise HTTPException(
            status_code=409,
            detail=(
                f"A throughput benchmark is already running "
                f"(bench_id={active.bench_id}, model_id={active.request.model_id})."
            ),
        )

    from .context_benchmark import get_active_run as get_active_context_run

    context_active = get_active_context_run()
    if context_active is not None:
        raise HTTPException(
            status_code=409,
            detail=(
                f"A context benchmark is already running "
                f"(bench_id={context_active.bench_id}, "
                f"model_id={context_active.request.model_id})."
            ),
        )

    body = await request.json()
    try:
        bench_request = BenchmarkRequest(**body)
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

    # Validate model exists and is an LLM. External runs target a remote
    # model — nothing to validate locally.
    if bench_request.external is None:
        entry = engine_pool.get_entry(bench_request.model_id)
        if entry is None:
            raise HTTPException(
                status_code=404, detail=f"Model not found: {bench_request.model_id}"
            )
        if entry.model_type not in ("llm", "vlm", None):
            raise HTTPException(
                status_code=400,
                detail=f"Model {bench_request.model_id} is not a supported model (type: {entry.model_type})",
            )

    # Cleanup old runs
    cleanup_old_runs()

    # Create and start the benchmark
    run = create_run(bench_request)
    total_tests = len(bench_request.prompt_lengths) + len(bench_request.batch_sizes) * 2

    run.task = asyncio.create_task(run_benchmark(run, engine_pool))

    logger.info(
        f"Benchmark started: {run.bench_id} model={bench_request.model_id} "
        f"tests={total_tests}"
    )

    return {
        "bench_id": run.bench_id,
        "status": "started",
        "total_tests": total_tests,
    }


@router.get("/api/bench/{bench_id}/stream")
async def stream_benchmark(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Stream benchmark progress via Server-Sent Events."""
    import json

    from fastapi.responses import StreamingResponse

    from .benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Benchmark not found: {bench_id}")

    async def event_generator():
        # Replay-then-attach: see /api/bench/accuracy/{id}/stream for the
        # full rationale. The bench stream's terminal events are
        # `upload_done` and `error` — `done` only marks the boundary
        # between tests and upload.
        seen = 0
        try:
            while True:
                async with run.cond:
                    while seen >= len(run.events) and not run.terminal:
                        try:
                            await asyncio.wait_for(run.cond.wait(), timeout=60.0)
                        except TimeoutError:
                            break
                    new = list(run.events[seen:])
                    seen = len(run.events)
                    done = run.terminal

                for ev in new:
                    yield f"data: {json.dumps(ev)}\n\n"
                if not new and not done:
                    yield ": keepalive\n\n"
                if done:
                    break
        except asyncio.CancelledError:
            pass

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@router.post("/api/bench/{bench_id}/cancel")
async def cancel_benchmark(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Cancel a running benchmark."""
    from .benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Benchmark not found: {bench_id}")

    if run.status != "running":
        raise HTTPException(
            status_code=400,
            detail=f"Benchmark is not running (status: {run.status})",
        )

    if run.task and not run.task.done():
        run.task.cancel()

    return {"status": "cancelled", "bench_id": bench_id}


@router.get("/api/bench/{bench_id}/results")
async def get_benchmark_results(
    bench_id: str,
    is_admin: bool = Depends(require_admin),
):
    """Get results from a completed benchmark."""
    from .benchmark import get_run

    run = get_run(bench_id)
    if run is None:
        raise HTTPException(status_code=404, detail=f"Benchmark not found: {bench_id}")

    return {
        "bench_id": run.bench_id,
        "status": run.status,
        "context_profile": run.request.context_profile.value,
        "results": run.results,
        "error": run.error_message if run.error_message else None,
        "upload_state": run.upload_state,
    }


@router.get("/api/device-info")
async def get_device_info(
    is_admin: bool = Depends(require_admin),
):
    """Get device hardware info and owner_hash for omlx.ai integration."""
    from ..utils.hardware import (
        compute_owner_hash,
        get_chip_name,
        get_gpu_core_count,
        get_io_platform_uuid,
        get_total_memory_gb,
        parse_chip_info,
    )

    chip_string = get_chip_name()
    chip_name, chip_variant = parse_chip_info(chip_string)
    memory_gb = round(get_total_memory_gb())
    gpu_cores = get_gpu_core_count()

    owner_hash = None
    io_uuid = get_io_platform_uuid()
    if io_uuid:
        full_hash = compute_owner_hash(io_uuid, chip_name, gpu_cores, memory_gb)
        owner_hash = full_hash[:-1]  # Strip verify character for URL

    return {
        "chip_name": chip_name,
        "chip_variant": chip_variant,
        "memory_gb": memory_gb,
        "gpu_cores": gpu_cores,
        "owner_hash": owner_hash,
    }


# =============================================================================
# Update Check
# =============================================================================

_update_cache: dict[str, dict[str, Any]] = {}
_update_cache_time: dict[str, float] = {}
_UPDATE_CACHE_TTL = 3600  # 1 hour
_UPDATE_PREFS_PATH = (
    Path.home() / "Library" / "Application Support" / "oMLX" / "update-prefs.json"
)


def _read_update_channel() -> str:
    try:
        data = json.loads(_UPDATE_PREFS_PATH.read_text())
    except Exception:
        return "stable"
    return normalize_update_channel(data.get("channel"))


@router.get("/api/update-check")
async def check_update(
    is_admin: bool = Depends(require_admin),
):
    """Check GitHub Releases for newer oMLX version (cached 24h)."""
    global _update_cache, _update_cache_time

    now = time.time()
    channel = _read_update_channel()

    if not isinstance(_update_cache, dict) or _update_cache is None:
        _update_cache = {}
    if not isinstance(_update_cache_time, dict) or _update_cache_time is None:
        _update_cache_time = {}

    cached = _update_cache.get(channel)
    cached_time = _update_cache_time.get(channel, 0.0)
    if cached is not None and now - cached_time < _UPDATE_CACHE_TTL:
        return cached

    no_update = {
        "update_available": False,
        "latest_version": None,
        "release_url": None,
        "update_channel": channel,
    }

    try:
        # Use the releases list (not /releases/latest) and filter by the
        # user's update channel. GitHub's prerelease flag has historically
        # been unreliable for rc/dev tags, so release_check validates tags too.
        resp = await asyncio.to_thread(
            requests.get,
            "https://api.github.com/repos/jundot/omlx/releases",
            params={"per_page": 20},
            timeout=5,
        )
        if resp.status_code != 200:
            _update_cache[channel] = no_update
            _update_cache_time[channel] = now
            return _update_cache[channel]

        data = select_latest_release(resp.json(), channel=channel)
        if data is None:
            _update_cache[channel] = no_update
            _update_cache_time[channel] = now
            return _update_cache[channel]

        latest = data["tag_name"].lstrip("v")

        try:
            from packaging.version import Version

            update_available = Version(latest) > Version(_omlx_version)
        except Exception:
            update_available = False

        if update_available:
            _update_cache[channel] = {
                "update_available": True,
                "latest_version": latest,
                "release_url": data.get("html_url"),
                "update_channel": channel,
            }
        else:
            _update_cache[channel] = no_update

        _update_cache_time[channel] = now
    except Exception:
        _update_cache[channel] = no_update
        _update_cache_time[channel] = now

    return _update_cache[channel]


# =============================================================================
# oQ Quantization API Routes
# =============================================================================


@router.get("/api/oq/models")
async def list_oq_models(is_admin: bool = Depends(require_admin)):
    """List non-quantized models available for oQ quantization."""
    if _oq_manager is None:
        raise HTTPException(status_code=503, detail="oQ quantizer not initialized")
    source_models, all_models = await _oq_manager.list_quantizable_models()
    return {"models": source_models, "all_models": all_models}


@router.get("/api/oq/estimate")
async def estimate_oq(
    model_path: str,
    oq_level: float,
    preserve_mtp: bool = False,
    is_admin: bool = Depends(require_admin),
):
    """Estimate effective bpw and output size for a model at given oQ level."""
    from ..oq import estimate_bpw_and_size

    try:
        result = await asyncio.to_thread(
            estimate_bpw_and_size,
            model_path,
            oq_level,
            64,  # group_size (default)
            preserve_mtp,
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.post("/api/oq/start")
async def start_oq_quantization(
    request: OQStartRequest,
    is_admin: bool = Depends(require_admin),
):
    """Start an oQ quantization task."""
    from ..oq import OQ_LEVELS

    if _oq_manager is None:
        raise HTTPException(status_code=503, detail="oQ quantizer not initialized")
    if request.oq_level not in OQ_LEVELS:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid oQ level. Must be one of {sorted(OQ_LEVELS)}",
        )
    if request.dtype not in ("bfloat16", "float16"):
        raise HTTPException(
            status_code=400,
            detail="Invalid dtype. Must be 'bfloat16' or 'float16'",
        )
    if request.enhanced:
        if not 1 <= request.imatrix_num_samples <= 4096:
            raise HTTPException(
                status_code=400,
                detail="Invalid imatrix_num_samples. Must be between 1 and 4096.",
            )
        if not 64 <= request.imatrix_seq_length <= 8192:
            raise HTTPException(
                status_code=400,
                detail="Invalid imatrix_seq_length. Must be between 64 and 8192.",
            )
    is_paro, _ = _paroquant_compat_for_model({"model_path": request.model_path})
    if is_paro:
        raise HTTPException(
            status_code=400,
            detail=(
                "Model is already quantized with paroquant; "
                "oQ re-quantization is not supported"
            ),
        )
    try:
        task = await _oq_manager.start_quantization(
            model_path=request.model_path,
            oq_level=request.oq_level,
            group_size=request.group_size,
            sensitivity_model_path=request.sensitivity_model_path,
            text_only=request.text_only,
            dtype=request.dtype,
            preserve_mtp=request.preserve_mtp,
            auto_proxy_sensitivity=request.auto_proxy_sensitivity,
            enhanced=request.enhanced,
            imatrix_cache_path=request.imatrix_cache_path,
            imatrix_reuse_cache=request.imatrix_reuse_cache,
            imatrix_strict=request.imatrix_strict,
            imatrix_num_samples=request.imatrix_num_samples,
            imatrix_seq_length=request.imatrix_seq_length,
            mtp_assistant_model_path=request.mtp_assistant_model_path,
        )
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.get("/api/oq/tasks")
async def list_oq_tasks(is_admin: bool = Depends(require_admin)):
    """List all quantization tasks."""
    if _oq_manager is None:
        raise HTTPException(status_code=503, detail="oQ quantizer not initialized")
    return {"tasks": _oq_manager.get_tasks()}


@router.post("/api/oq/cancel/{task_id}")
async def cancel_oq_task(task_id: str, is_admin: bool = Depends(require_admin)):
    """Cancel an active quantization task."""
    if _oq_manager is None:
        raise HTTPException(status_code=503, detail="oQ quantizer not initialized")
    success = await _oq_manager.cancel_quantization(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or not cancellable")
    return {"success": True}


@router.delete("/api/oq/task/{task_id}")
async def remove_oq_task(task_id: str, is_admin: bool = Depends(require_admin)):
    """Remove a completed/failed/cancelled task."""
    if _oq_manager is None:
        raise HTTPException(status_code=503, detail="oQ quantizer not initialized")
    success = _oq_manager.remove_task(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or still active")
    return {"success": True}


# =============================================================================
# HuggingFace Upload Endpoints
# =============================================================================


@router.post("/api/upload/validate-token")
async def validate_upload_token(
    request: HFValidateTokenRequest,
    is_admin: bool = Depends(require_admin),
):
    """Validate a HuggingFace token and return user info."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    try:
        result = await _hf_uploader.validate_token(request.hf_token)
        return result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.get("/api/upload/oq-models")
async def list_upload_oq_models(is_admin: bool = Depends(require_admin)):
    """List local oQ models available for upload."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    oq_models = await _hf_uploader.list_oq_models()
    all_models = await _hf_uploader.list_all_models()
    return {"oq_models": oq_models, "all_models": all_models}


@router.post("/api/upload/start")
async def start_upload(
    request: HFUploadRequest,
    is_admin: bool = Depends(require_admin),
):
    """Start an upload task to HuggingFace Hub."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    try:
        task = await _hf_uploader.start_upload(
            model_path=request.model_path,
            repo_id=request.repo_id,
            token=request.hf_token,
            readme_source_path=request.readme_source_path,
            auto_readme=request.auto_readme,
            redownload_notice=request.redownload_notice,
            private=request.private,
        )
        return {"success": True, "task": task.to_dict()}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.get("/api/upload/tasks")
async def list_upload_tasks(is_admin: bool = Depends(require_admin)):
    """List all upload tasks."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    return {"tasks": _hf_uploader.get_tasks()}


@router.post("/api/upload/cancel/{task_id}")
async def cancel_upload_task(task_id: str, is_admin: bool = Depends(require_admin)):
    """Cancel an active or pending upload task."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    success = await _hf_uploader.cancel_upload(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or not cancellable")
    return {"success": True}


@router.delete("/api/upload/task/{task_id}")
async def remove_upload_task(task_id: str, is_admin: bool = Depends(require_admin)):
    """Remove a completed/failed/cancelled upload task."""
    if _hf_uploader is None:
        raise HTTPException(status_code=503, detail="HF Uploader not initialized")
    success = _hf_uploader.remove_task(task_id)
    if not success:
        raise HTTPException(status_code=404, detail="Task not found or still active")
    return {"success": True}


# ===== FILE: omlx/admin/vendor_deps.py =====

#!/usr/bin/env python3
"""Download vendored dependencies for offline admin panel.

All libraries use permissive licenses (MIT/ISC/BSD/OFL) that allow bundling.
Run this script to download/update all CDN dependencies to static/.

Usage:
    cd omlx/omlx/admin
    python vendor_deps.py
"""

import re
import ssl
import urllib.request
from pathlib import Path

STATIC = Path(__file__).parent / "static"

# SSL context for HTTPS downloads
SSL_CTX = ssl.create_default_context()


def _download(url: str, dest: Path, description: str = "", optional: bool = False) -> bool:
    """Download a file from URL to destination path.

    Args:
        optional: If True, silently skip 404 errors (some font variants don't exist).

    Returns:
        True if downloaded or already exists, False if skipped.
    """
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists():
        print(f"  [skip] {dest.name} (already exists)")
        return True
    label = description or dest.name
    print(f"  [download] {label} <- {url}")
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    try:
        with urllib.request.urlopen(req, context=SSL_CTX) as resp:
            dest.write_bytes(resp.read())
        return True
    except urllib.error.HTTPError as e:
        if optional and e.code == 404:
            print(f"  [skip] {dest.name} (not available)")
            return False
        raise


# =========================================================================
# JavaScript dependencies
# =========================================================================
JS_DEPS = {
    # Alpine.js 3.14.8 (MIT)
    "js/alpine.min.js": "https://cdn.jsdelivr.net/npm/alpinejs@3.14.8/dist/cdn.min.js",
    # Lucide Icons 0.453.0 (ISC)
    "js/lucide.min.js": "https://unpkg.com/lucide@0.453.0/dist/umd/lucide.min.js",
    # Marked 12.0.0 (MIT)
    "js/marked.umd.js": "https://cdn.jsdelivr.net/npm/marked@12.0.0/lib/marked.umd.js",
    # marked-highlight 2.0.6 (MIT)
    "js/marked-highlight.umd.js": "https://cdn.jsdelivr.net/npm/marked-highlight@2.0.6/lib/index.umd.js",
    # Highlight.js 11.9.0 core (BSD-3-Clause)
    "js/highlight.min.js": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js",
    # Highlight.js language packs
    "js/hljs-python.min.js": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/languages/python.min.js",
    "js/hljs-javascript.min.js": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/languages/javascript.min.js",
    "js/hljs-bash.min.js": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/languages/bash.min.js",
    "js/hljs-json.min.js": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/languages/json.min.js",
    # KaTeX 0.16.9 (MIT)
    "js/katex.min.js": "https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.js",
    "js/katex-auto-render.min.js": "https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/contrib/auto-render.min.js",
}

# =========================================================================
# CSS dependencies
# =========================================================================
CSS_DEPS = {
    # Highlight.js themes (BSD-3-Clause)
    "css/hljs-github.min.css": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github.min.css",
    "css/hljs-github-dark.min.css": "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github-dark.min.css",
    # KaTeX CSS (MIT) - references fonts/ relative path
    "css/katex.min.css": "https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.css",
}


def download_js_css() -> None:
    """Download JavaScript and CSS dependencies."""
    print("\n=== JavaScript Dependencies ===")
    for dest_rel, url in JS_DEPS.items():
        _download(url, STATIC / dest_rel)

    print("\n=== CSS Dependencies ===")
    for dest_rel, url in CSS_DEPS.items():
        _download(url, STATIC / dest_rel)


# =========================================================================
# KaTeX fonts
# =========================================================================
KATEX_VERSION = "0.16.9"
KATEX_FONT_BASE = f"https://cdn.jsdelivr.net/npm/katex@{KATEX_VERSION}/dist/fonts"

# All KaTeX font files referenced in katex.min.css
KATEX_FONTS = [
    "KaTeX_AMS-Regular",
    "KaTeX_Caligraphic-Bold",
    "KaTeX_Caligraphic-Regular",
    "KaTeX_Fraktur-Bold",
    "KaTeX_Fraktur-Regular",
    "KaTeX_Main-Bold",
    "KaTeX_Main-BoldItalic",
    "KaTeX_Main-Italic",
    "KaTeX_Main-Regular",
    "KaTeX_Math-BoldItalic",
    "KaTeX_Math-Italic",
    "KaTeX_SansSerif-Bold",
    "KaTeX_SansSerif-Italic",
    "KaTeX_SansSerif-Regular",
    "KaTeX_Script-Regular",
    "KaTeX_Size1-Regular",
    "KaTeX_Size2-Regular",
    "KaTeX_Size3-Regular",
    "KaTeX_Size4-Regular",
    "KaTeX_Typewriter-Regular",
]


def download_katex_fonts() -> None:
    """Download KaTeX font files (woff2 + ttf fallback)."""
    print("\n=== KaTeX Fonts ===")
    # Place in css/fonts/ so katex.min.css relative path works (url(fonts/...))
    fonts_dir = STATIC / "css" / "fonts"
    for font_name in KATEX_FONTS:
        for ext in ("woff2", "ttf"):
            url = f"{KATEX_FONT_BASE}/{font_name}.{ext}"
            _download(url, fonts_dir / f"{font_name}.{ext}", optional=True)


# =========================================================================
# Inter font (SIL Open Font License)
# =========================================================================
INTER_WEIGHTS = [300, 400, 500, 600, 700, 800]
INTER_FONT_BASE = "https://cdn.jsdelivr.net/fontsource/fonts/inter@latest"


def download_inter_fonts() -> None:
    """Download Inter font files and create @font-face CSS."""
    print("\n=== Inter Font ===")
    inter_dir = STATIC / "fonts" / "inter"
    inter_dir.mkdir(parents=True, exist_ok=True)

    for weight in INTER_WEIGHTS:
        url = f"{INTER_FONT_BASE}/latin-{weight}-normal.woff2"
        _download(url, inter_dir / f"inter-latin-{weight}-normal.woff2")

    # Generate @font-face CSS
    css_path = inter_dir / "inter.css"
    if css_path.exists():
        print("  [skip] inter.css (already exists)")
        return

    print("  [generate] inter.css")
    css_parts = []
    for weight in INTER_WEIGHTS:
        css_parts.append(f"""@font-face {{
  font-family: 'Inter';
  font-style: normal;
  font-weight: {weight};
  font-display: swap;
  src: url('./inter-latin-{weight}-normal.woff2') format('woff2');
  unicode-range: U+0000-00FF, U+0131, U+0152-0153, U+02BB-02BC, U+02C6, U+02DA,
    U+02DC, U+0304, U+0308, U+0329, U+2000-206F, U+20AC, U+2122, U+2191, U+2193,
    U+2212, U+2215, U+FEFF, U+FFFD;
}}""")
    css_path.write_text("\n\n".join(css_parts) + "\n")


# =========================================================================
# CJK fonts (SIL Open Font License)
# =========================================================================
CJK_FONTS = {
    # (font_family, fontsource_id, subset, dir_name, file_prefix)
    "noto-sans-sc": ("Noto Sans SC", "noto-sans-sc", "chinese-simplified", "NotoSansSC"),
    "noto-sans-tc": ("Noto Sans TC", "noto-sans-tc", "chinese-traditional", "NotoSansTC"),
    "noto-sans-kr": ("Noto Sans KR", "noto-sans-kr", "korean", "NotoSansKR"),
    "noto-sans-jp": ("Noto Sans JP", "noto-sans-jp", "japanese", "NotoSansJP"),
}
CJK_WEIGHTS = {400: "Regular", 500: "Medium", 700: "Bold"}
CJK_FONT_BASE = "https://cdn.jsdelivr.net/fontsource/fonts"


def download_cjk_fonts() -> None:
    """Download CJK font files (Noto Sans SC/TC/KR/JP) and create @font-face CSS."""
    print("\n=== CJK Fonts ===")
    for dir_name, (family, fontsource_id, subset, prefix) in CJK_FONTS.items():
        font_dir = STATIC / "fonts" / dir_name
        font_dir.mkdir(parents=True, exist_ok=True)

        for weight, weight_name in CJK_WEIGHTS.items():
            filename = f"{prefix}-{weight_name}.woff2"
            url = f"{CJK_FONT_BASE}/{fontsource_id}@latest/{subset}-{weight}-normal.woff2"
            _download(url, font_dir / filename)

        # Generate @font-face CSS
        css_path = font_dir / f"{dir_name}.css"
        if css_path.exists():
            print(f"  [skip] {dir_name}.css (already exists)")
            continue

        print(f"  [generate] {dir_name}.css")
        comment = {
            "noto-sans-sc": "Simplified Chinese",
            "noto-sans-tc": "Traditional Chinese",
            "noto-sans-kr": "Korean",
            "noto-sans-jp": "Japanese",
        }[dir_name]
        css_parts = [f"/* {family} - {comment} */"]
        for weight, weight_name in CJK_WEIGHTS.items():
            css_parts.append(f"""@font-face {{
  font-family: '{family}';
  font-style: normal;
  font-weight: {weight};
  font-display: swap;
  src: url('{prefix}-{weight_name}.woff2') format('woff2');
}}""")
        css_path.write_text("\n".join(css_parts) + "\n")


def main() -> None:
    print(f"Vendor directory: {STATIC}")
    download_js_css()
    download_katex_fonts()
    download_inter_fonts()
    download_cjk_fonts()
    print("\n=== Done! ===")

    # Summary
    total = 0
    for p in STATIC.rglob("*"):
        if p.is_file() and p.suffix != ".svg":
            total += p.stat().st_size
    print(f"Total vendored size: {total / 1024 / 1024:.1f} MB")


if __name__ == "__main__":
    main()


# ===== FILE: omlx/api/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""
API models, utilities, and tool calling support for oMLX.

This module provides shared components used by the server:
- Pydantic models for OpenAI-compatible API
- Pydantic models for Anthropic Messages API
- Utility functions for text processing
- Tool calling parsing and conversion
"""

from .openai_models import (
    # Content types
    ContentPart,
    FileContent,
    Message,
    # Tool calling
    FunctionCall,
    ToolCall,
    ToolDefinition,
    # Structured output
    ResponseFormat,
    ResponseFormatJsonSchema,
    # Chat requests/responses
    ChatCompletionRequest,
    ChatCompletionChoice,
    ChatCompletionResponse,
    AssistantMessage,
    # Completion requests/responses
    CompletionRequest,
    CompletionChoice,
    CompletionResponse,
    # Common
    Usage,
    ModelInfo,
    ModelsResponse,
    # MCP
    MCPToolInfo,
    MCPToolsResponse,
    MCPServerInfo,
    MCPServersResponse,
    MCPExecuteRequest,
    MCPExecuteResponse,
)

from .utils import (
    clean_output_text,
    clean_special_tokens,
    extract_text_content,
    SPECIAL_TOKENS_PATTERN,
)

from .thinking import (
    ThinkingParser,
    extract_thinking,
)

from .tool_calling import (
    parse_tool_calls,
    convert_tools_for_template,
    # Structured output
    parse_json_output,
    validate_json_schema,
    extract_json_from_text,
    build_json_system_prompt,
)

from .anthropic_models import (
    # Content blocks
    ContentBlockText,
    ContentBlockImage,
    ContentBlockToolUse,
    ContentBlockToolResult,
    ContentBlockInputAudio,
    # Messages
    SystemContent,
    AnthropicMessage,
    AnthropicTool,
    ToolChoice,
    ThinkingConfig,
    # Request/Response
    MessagesRequest as AnthropicMessagesRequest,
    MessagesResponse as AnthropicMessagesResponse,
    AnthropicUsage,
    # Token counting
    TokenCountRequest,
    TokenCountResponse,
    # Streaming events
    MessageStartEvent,
    ContentBlockStartEvent,
    ContentBlockDeltaEvent,
    ContentBlockStopEvent,
    MessageDeltaEvent,
    MessageStopEvent,
    PingEvent,
    ErrorEvent,
    TextDelta,
    InputJsonDelta,
    # Error response
    AnthropicErrorResponse,
    AnthropicErrorDetail,
)

from .anthropic_utils import (
    convert_anthropic_to_internal,
    convert_anthropic_tools_to_internal,
    convert_internal_to_anthropic_response,
    map_finish_reason_to_stop_reason,
    format_sse_event,
    create_message_start_event,
    create_content_block_start_event,
    create_text_delta_event,
    create_input_json_delta_event,
    create_content_block_stop_event,
    create_message_delta_event,
    create_message_stop_event,
    create_ping_event,
    create_error_event,
)

from .embedding_models import (
    EmbeddingRequest,
    EmbeddingResponse,
    EmbeddingData,
    EmbeddingUsage,
)

from .embedding_utils import (
    encode_embedding_base64,
    truncate_embedding,
    count_tokens,
    normalize_input,
)

# MCP routes
from .mcp_routes import router as mcp_router, set_mcp_manager_getter

__all__ = [
    # Models
    "ContentPart",
    "FileContent",
    "Message",
    "FunctionCall",
    "ToolCall",
    "ToolDefinition",
    "ResponseFormat",
    "ResponseFormatJsonSchema",
    "ChatCompletionRequest",
    "ChatCompletionChoice",
    "ChatCompletionResponse",
    "AssistantMessage",
    "CompletionRequest",
    "CompletionChoice",
    "CompletionResponse",
    "Usage",
    "ModelInfo",
    "ModelsResponse",
    "MCPToolInfo",
    "MCPToolsResponse",
    "MCPServerInfo",
    "MCPServersResponse",
    "MCPExecuteRequest",
    "MCPExecuteResponse",
    # Utils
    "clean_output_text",
    "clean_special_tokens",
    "extract_text_content",
    "SPECIAL_TOKENS_PATTERN",
    # Thinking
    "ThinkingParser",
    "extract_thinking",
    # Tool calling
    "parse_tool_calls",
    "convert_tools_for_template",
    # Structured output
    "parse_json_output",
    "validate_json_schema",
    "extract_json_from_text",
    "build_json_system_prompt",
    # Anthropic models
    "ContentBlockText",
    "ContentBlockImage",
    "ContentBlockToolUse",
    "ContentBlockToolResult",
    "SystemContent",
    "AnthropicMessage",
    "AnthropicTool",
    "ToolChoice",
    "ThinkingConfig",
    "AnthropicMessagesRequest",
    "AnthropicMessagesResponse",
    "AnthropicUsage",
    "TokenCountRequest",
    "TokenCountResponse",
    "MessageStartEvent",
    "ContentBlockStartEvent",
    "ContentBlockDeltaEvent",
    "ContentBlockStopEvent",
    "MessageDeltaEvent",
    "MessageStopEvent",
    "PingEvent",
    "ErrorEvent",
    "TextDelta",
    "InputJsonDelta",
    "AnthropicErrorResponse",
    "AnthropicErrorDetail",
    # Anthropic utils
    "convert_anthropic_to_internal",
    "convert_anthropic_tools_to_internal",
    "convert_internal_to_anthropic_response",
    "map_finish_reason_to_stop_reason",
    "format_sse_event",
    "create_message_start_event",
    "create_content_block_start_event",
    "create_text_delta_event",
    "create_input_json_delta_event",
    "create_content_block_stop_event",
    "create_message_delta_event",
    "create_message_stop_event",
    "create_ping_event",
    "create_error_event",
    # Embedding models
    "EmbeddingRequest",
    "EmbeddingResponse",
    "EmbeddingData",
    "EmbeddingUsage",
    # Embedding utils
    "encode_embedding_base64",
    "truncate_embedding",
    "count_tokens",
    "normalize_input",
    # MCP routes
    "mcp_router",
    "set_mcp_manager_getter",
]


# ===== FILE: omlx/api/adapters/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""
API Adapters for oMLX.

This package provides adapters for different API formats (OpenAI, Anthropic),
enabling clean separation between API-specific logic and core inference.
"""

from .base import (
    BaseAdapter,
    InternalMessage,
    InternalRequest,
    InternalResponse,
    StreamChunk,
)
from .openai import OpenAIAdapter
from .anthropic import AnthropicAdapter
from .sse_formatter import (
    SSEFormatter,
    OpenAISSEFormatter,
    AnthropicSSEFormatter,
)

__all__ = [
    # Base classes and types
    "BaseAdapter",
    "InternalMessage",
    "InternalRequest",
    "InternalResponse",
    "StreamChunk",
    # Adapters
    "OpenAIAdapter",
    "AnthropicAdapter",
    # SSE Formatters
    "SSEFormatter",
    "OpenAISSEFormatter",
    "AnthropicSSEFormatter",
]


# ===== FILE: omlx/api/adapters/anthropic.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Anthropic API adapter for oMLX.

This adapter handles conversion between Anthropic Messages API format and the
internal request/response format used by the inference engine.
"""

import json
import uuid
from typing import Any, List, Optional

from .base import (
    BaseAdapter,
    InternalMessage,
    InternalRequest,
    InternalResponse,
    StreamChunk,
)
from ..anthropic_models import (
    MessagesRequest as AnthropicMessagesRequest,
)
from ..anthropic_utils import (
    convert_anthropic_to_internal,
    convert_anthropic_tools_to_internal,
    convert_internal_to_anthropic_response,
    create_content_block_start_event,
    create_content_block_stop_event,
    create_error_event,
    create_input_json_delta_event,
    create_message_delta_event,
    create_message_start_event,
    create_message_stop_event,
    create_text_delta_event,
    map_finish_reason_to_stop_reason,
)


class AnthropicAdapter(BaseAdapter):
    """
    Adapter for Anthropic Messages API format.

    Handles conversion between Anthropic message requests/responses
    and the internal format used by the inference engine.
    """

    @property
    def name(self) -> str:
        return "anthropic"

    def parse_request(self, request: AnthropicMessagesRequest) -> InternalRequest:
        """
        Convert an Anthropic MessagesRequest to internal format.

        Args:
            request: Anthropic messages request.

        Returns:
            InternalRequest in unified format.
        """
        # Convert messages to internal format (includes system message handling)
        messages = convert_anthropic_to_internal(request)

        # Convert to InternalMessage objects
        internal_messages = []
        for msg in messages:
            internal_messages.append(
                InternalMessage(
                    role=msg.get("role", "user"),
                    content=msg.get("content", ""),
                )
            )

        # Convert tools if provided
        tools = None
        if request.tools:
            tools = convert_anthropic_tools_to_internal(request.tools)

        return InternalRequest(
            messages=internal_messages,
            max_tokens=request.max_tokens,
            temperature=request.temperature if request.temperature is not None else 1.0,
            top_p=request.top_p if request.top_p is not None else 1.0,
            top_k=request.top_k if request.top_k is not None else 0,
            stream=request.stream or False,
            stop=request.stop_sequences,
            tools=tools,
            model=request.model,
            request_id=f"msg_{uuid.uuid4().hex[:24]}",
        )

    def format_response(
        self,
        response: InternalResponse,
        request: AnthropicMessagesRequest,
    ) -> dict:
        """
        Convert an internal response to Anthropic Messages format.

        Args:
            response: Internal response object.
            request: Original Anthropic request.

        Returns:
            Response dict in Anthropic format.
        """
        return convert_internal_to_anthropic_response(
            text=response.text,
            finish_reason=response.finish_reason,
            prompt_tokens=response.prompt_tokens,
            completion_tokens=response.completion_tokens,
            model=request.model,
            tool_calls=response.tool_calls,
        )

    def format_stream_chunk(
        self,
        chunk: StreamChunk,
        request: AnthropicMessagesRequest,
    ) -> str:
        """
        Format a streaming chunk for SSE output in Anthropic format.

        Args:
            chunk: The stream chunk to format.
            request: Original Anthropic request.

        Returns:
            SSE-formatted string.
        """
        events = []

        # First chunk: send message_start and content_block_start
        if chunk.is_first:
            message_id = f"msg_{uuid.uuid4().hex[:24]}"
            events.append(create_message_start_event(message_id, request.model))
            events.append(create_content_block_start_event(0, "text"))

        # Text delta
        if chunk.text:
            events.append(create_text_delta_event(0, chunk.text))

        # Tool call delta
        if chunk.tool_call_delta:
            partial_json = json.dumps(chunk.tool_call_delta)
            events.append(create_input_json_delta_event(0, partial_json))

        # Last chunk: send content_block_stop, message_delta, message_stop
        if chunk.is_last:
            events.append(create_content_block_stop_event(0))
            events.append(
                create_message_delta_event(
                    stop_reason=map_finish_reason_to_stop_reason(
                        chunk.finish_reason, bool(chunk.tool_call_delta)
                    ),
                    output_tokens=chunk.completion_tokens,
                )
            )
            events.append(create_message_stop_event())

        return "".join(events)

    def format_stream_end(self, request: AnthropicMessagesRequest) -> str:
        """
        Format the stream end marker for Anthropic format.

        Anthropic doesn't use [DONE] - the message_stop event is the end.

        Args:
            request: Original Anthropic request.

        Returns:
            Empty string (no additional marker needed).
        """
        return ""

    def create_error_response(
        self,
        error: str,
        error_type: str = "api_error",
        status_code: int = 500,
    ) -> dict:
        """
        Create an error response in Anthropic format.

        Args:
            error: Error message.
            error_type: Type of error (e.g., "invalid_request_error").
            status_code: HTTP status code (not used in Anthropic format).

        Returns:
            Error response dict in Anthropic format.
        """
        return {
            "type": "error",
            "error": {
                "type": error_type,
                "message": error,
            },
        }

    def format_error_event(
        self,
        error: str,
        error_type: str = "api_error",
    ) -> str:
        """
        Create an error event in Anthropic SSE format.

        Args:
            error: Error message.
            error_type: Type of error.

        Returns:
            SSE-formatted error event.
        """
        return create_error_event(error_type, error)


# ===== FILE: omlx/api/adapters/base.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Base adapter interface for API format conversion.

This module defines the abstract interface that all API adapters must implement,
plus internal data structures for request/response handling.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, List, Optional, Union


@dataclass
class InternalMessage:
    """Internal representation of a chat message."""

    role: str
    content: str
    name: Optional[str] = None
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_call_id: Optional[str] = None


@dataclass
class InternalRequest:
    """
    Internal request format used by the inference engine.

    This provides a unified format that all adapters convert to/from.
    """

    # Required fields
    messages: List[InternalMessage]

    # Generation parameters
    max_tokens: int = 2048
    temperature: float = 1.0
    top_p: float = 1.0
    top_k: int = 0
    min_p: float = 0.0
    presence_penalty: float = 0.0
    frequency_penalty: float = 0.0
    stream: bool = False

    # Stop conditions
    stop: Optional[List[str]] = None
    stop_token_ids: Optional[List[int]] = None

    # Tool calling
    tools: Optional[List[Dict[str, Any]]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None

    # Response format
    response_format: Optional[Dict[str, Any]] = None

    # Model
    model: Optional[str] = None

    # Metadata
    request_id: Optional[str] = None


@dataclass
class InternalResponse:
    """
    Internal response format from the inference engine.

    This provides a unified format that all adapters convert from.
    """

    # Generated content
    text: str
    finish_reason: Optional[str] = None
    reasoning_content: Optional[str] = None

    # Token counts
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cached_tokens: int = 0

    # Tool calls (parsed)
    tool_calls: Optional[List[Dict[str, Any]]] = None

    # Metadata
    request_id: Optional[str] = None
    model: Optional[str] = None


@dataclass
class StreamChunk:
    """A single chunk in a streaming response."""

    text: str = ""
    reasoning_content: Optional[str] = None
    finish_reason: Optional[str] = None
    tool_call_delta: Optional[Dict[str, Any]] = None
    is_first: bool = False
    is_last: bool = False

    # Token counts (usually only on last chunk)
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cached_tokens: int = 0


class BaseAdapter(ABC):
    """
    Abstract base class for API adapters.

    Adapters handle conversion between external API formats (OpenAI, Anthropic)
    and the internal request/response format used by the inference engine.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Return the adapter name (e.g., 'openai', 'anthropic')."""
        pass

    @abstractmethod
    def parse_request(self, request: Any) -> InternalRequest:
        """
        Convert an external API request to internal format.

        Args:
            request: The external API request object.

        Returns:
            InternalRequest in unified format.
        """
        pass

    @abstractmethod
    def format_response(
        self,
        response: InternalResponse,
        request: Any,
    ) -> Any:
        """
        Convert an internal response to external API format.

        Args:
            response: The internal response object.
            request: The original external request (for context).

        Returns:
            Response in the external API format.
        """
        pass

    @abstractmethod
    def format_stream_chunk(
        self,
        chunk: StreamChunk,
        request: Any,
    ) -> str:
        """
        Format a streaming chunk for SSE output.

        Args:
            chunk: The stream chunk to format.
            request: The original external request (for context).

        Returns:
            SSE-formatted string.
        """
        pass

    @abstractmethod
    def format_stream_end(self, request: Any) -> str:
        """
        Format the stream end marker.

        Args:
            request: The original external request (for context).

        Returns:
            SSE-formatted end marker.
        """
        pass

    @abstractmethod
    def create_error_response(
        self,
        error: str,
        error_type: str = "server_error",
        status_code: int = 500,
    ) -> dict:
        """
        Create an error response in the adapter's format.

        Args:
            error: Error message.
            error_type: Type of error (e.g., "invalid_request_error").
            status_code: HTTP status code.

        Returns:
            Error response dict in the adapter's format.
        """
        pass


# ===== FILE: omlx/api/adapters/openai.py =====

# SPDX-License-Identifier: Apache-2.0
"""
OpenAI API adapter for oMLX.

This adapter handles conversion between OpenAI API format and the internal
request/response format used by the inference engine.
"""

import json
import time
import uuid
from typing import Any, List, Optional

from .base import (
    BaseAdapter,
    InternalMessage,
    InternalRequest,
    InternalResponse,
    StreamChunk,
)
from ..openai_models import (
    AssistantMessage,
    ChatCompletionChoice,
    ChatCompletionChunk,
    ChatCompletionChunkChoice,
    ChatCompletionChunkDelta,
    ChatCompletionRequest,
    ChatCompletionResponse,
    PromptTokensDetails,
    Usage,
)
from ..thinking import extract_thinking
from ..utils import clean_special_tokens, extract_text_content
from ..tool_calling import convert_tools_for_template


class OpenAIAdapter(BaseAdapter):
    """
    Adapter for OpenAI API format.

    Handles conversion between OpenAI chat completion requests/responses
    and the internal format used by the inference engine.
    """

    @property
    def name(self) -> str:
        return "openai"

    def parse_request(self, request: ChatCompletionRequest) -> InternalRequest:
        """
        Convert an OpenAI ChatCompletionRequest to internal format.

        Args:
            request: OpenAI chat completion request.

        Returns:
            InternalRequest in unified format.
        """
        # Extract text content from messages
        messages = extract_text_content(request.messages)

        # Convert to internal messages
        internal_messages = [
            InternalMessage(
                role=msg.get("role", "user"),
                content=msg.get("content", ""),
            )
            for msg in messages
        ]

        # Convert tools if provided
        tools = None
        if request.tools:
            tools = convert_tools_for_template(request.tools)

        return InternalRequest(
            messages=internal_messages,
            max_tokens=request.max_tokens or 2048,
            temperature=request.temperature if request.temperature is not None else 1.0,
            top_p=request.top_p if request.top_p is not None else 1.0,
            min_p=request.min_p if request.min_p is not None else 0.0,
            presence_penalty=request.presence_penalty if request.presence_penalty is not None else 0.0,
            frequency_penalty=request.frequency_penalty if request.frequency_penalty is not None else 0.0,
            stream=request.stream or False,
            stop=request.stop if isinstance(request.stop, list) else (
                [request.stop] if request.stop else None
            ),
            tools=tools,
            tool_choice=request.tool_choice,
            response_format=request.response_format,
            model=request.model,
            request_id=f"chatcmpl-{uuid.uuid4().hex[:12]}",
        )

    def format_response(
        self,
        response: InternalResponse,
        request: ChatCompletionRequest,
    ) -> ChatCompletionResponse:
        """
        Convert an internal response to OpenAI ChatCompletionResponse.

        Args:
            response: Internal response object.
            request: Original OpenAI request.

        Returns:
            ChatCompletionResponse in OpenAI format.
        """
        # Separate thinking from content
        raw_text = clean_special_tokens(response.text) if response.text else ""
        thinking_content, regular_content = extract_thinking(raw_text)
        content = regular_content.strip() if regular_content else None

        # Determine finish reason
        finish_reason = (
            "tool_calls" if response.tool_calls else response.finish_reason
        )

        return ChatCompletionResponse(
            id=response.request_id or f"chatcmpl-{uuid.uuid4().hex[:12]}",
            model=request.model,
            choices=[
                ChatCompletionChoice(
                    message=AssistantMessage(
                        content=content,
                        reasoning_content=thinking_content if thinking_content else None,
                        tool_calls=response.tool_calls,
                    ),
                    finish_reason=finish_reason,
                )
            ],
            usage=Usage(
                prompt_tokens=response.prompt_tokens,
                completion_tokens=response.completion_tokens,
                total_tokens=response.prompt_tokens + response.completion_tokens,
                prompt_tokens_details=PromptTokensDetails(
                    cached_tokens=response.cached_tokens,
                ),
            ),
        )

    def format_stream_chunk(
        self,
        chunk: StreamChunk,
        request: ChatCompletionRequest,
    ) -> str:
        """
        Format a streaming chunk for SSE output in OpenAI format.

        Args:
            chunk: The stream chunk to format.
            request: Original OpenAI request.

        Returns:
            SSE-formatted string.
        """
        request_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"

        delta = ChatCompletionChunkDelta(
            content=chunk.text if chunk.text else None,
            reasoning_content=chunk.reasoning_content if chunk.reasoning_content else None,
            tool_calls=chunk.tool_call_delta,
        )

        # Add role on first chunk
        if chunk.is_first:
            delta.role = "assistant"

        response = ChatCompletionChunk(
            id=request_id,
            model=request.model,
            choices=[
                ChatCompletionChunkChoice(
                    delta=delta,
                    finish_reason=chunk.finish_reason,
                )
            ],
        )

        # Add usage on last chunk if available
        if chunk.is_last and (chunk.prompt_tokens > 0 or chunk.completion_tokens > 0):
            response.usage = Usage(
                prompt_tokens=chunk.prompt_tokens,
                completion_tokens=chunk.completion_tokens,
                total_tokens=chunk.prompt_tokens + chunk.completion_tokens,
                prompt_tokens_details=PromptTokensDetails(
                    cached_tokens=chunk.cached_tokens,
                ),
            )

        return f"data: {response.model_dump_json(exclude_none=True)}\n\n"

    def format_stream_end(self, request: ChatCompletionRequest) -> str:
        """
        Format the stream end marker for OpenAI format.

        Args:
            request: Original OpenAI request.

        Returns:
            SSE-formatted end marker.
        """
        return "data: [DONE]\n\n"

    def create_error_response(
        self,
        error: str,
        error_type: str = "server_error",
        status_code: int = 500,
    ) -> dict:
        """
        Create an error response in OpenAI format.

        Args:
            error: Error message.
            error_type: Type of error.
            status_code: HTTP status code.

        Returns:
            Error response dict.
        """
        return {
            "error": {
                "message": error,
                "type": error_type,
                "param": None,
                "code": status_code,
            }
        }


# ===== FILE: omlx/api/adapters/sse_formatter.py =====

# SPDX-License-Identifier: Apache-2.0
"""
SSE (Server-Sent Events) formatting utilities.

This module provides abstract and concrete implementations for formatting
SSE events in different API formats (OpenAI, Anthropic).
"""

import json
from abc import ABC, abstractmethod
from typing import Any, Dict


class SSEFormatter(ABC):
    """Abstract base class for SSE event formatting."""

    @abstractmethod
    def format_event(self, event_type: str, data: Dict[str, Any]) -> str:
        """
        Format an SSE event.

        Args:
            event_type: Type of event (e.g., "content_block_delta").
            data: Event data to serialize.

        Returns:
            Formatted SSE event string.
        """
        pass

    @abstractmethod
    def format_end(self) -> str:
        """
        Format the stream end marker.

        Returns:
            Formatted SSE end marker string.
        """
        pass


class OpenAISSEFormatter(SSEFormatter):
    """
    SSE formatter for OpenAI API format.

    OpenAI uses a simple format: `data: {json}\n\n`
    The stream ends with `data: [DONE]\n\n`
    """

    def format_event(self, event_type: str, data: Dict[str, Any]) -> str:
        """
        Format an SSE event in OpenAI format.

        Note: OpenAI doesn't use event types in SSE, only data lines.

        Args:
            event_type: Ignored in OpenAI format.
            data: Event data to serialize.

        Returns:
            Formatted SSE event string.
        """
        return f"data: {json.dumps(data)}\n\n"

    def format_end(self) -> str:
        """
        Format the stream end marker for OpenAI format.

        Returns:
            The [DONE] marker.
        """
        return "data: [DONE]\n\n"


class AnthropicSSEFormatter(SSEFormatter):
    """
    SSE formatter for Anthropic API format.

    Anthropic uses event types: `event: {type}\ndata: {json}\n\n`
    The stream ends with a message_stop event (no separate [DONE] marker).
    """

    def format_event(self, event_type: str, data: Dict[str, Any]) -> str:
        """
        Format an SSE event in Anthropic format.

        Args:
            event_type: Type of event (e.g., "message_start", "content_block_delta").
            data: Event data to serialize.

        Returns:
            Formatted SSE event string with event type.
        """
        return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"

    def format_end(self) -> str:
        """
        Format the stream end marker for Anthropic format.

        Anthropic doesn't use a separate end marker - the message_stop
        event is the end of the stream.

        Returns:
            Empty string.
        """
        return ""


# ===== FILE: omlx/api/anthropic_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Pydantic models for Anthropic Messages API.

These models define the request and response schemas for:
- Anthropic Messages API (/v1/messages)
- Streaming events
- Tool calling in Anthropic format
"""

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from omlx.api.shared_models import IDPrefix, generate_id

# =============================================================================
# Content Blocks
# =============================================================================


class ContentBlockText(BaseModel):
    """Text content block."""

    type: Literal["text"] = "text"
    text: str
    cache_control: dict[str, str] | None = None


class ContentBlockImage(BaseModel):
    """Image content block with source data."""

    type: Literal["image"] = "image"
    source: dict[str, Any]  # {"type": "base64"|"url", "media_type": "...", "data"|"url": "..."}


class ContentBlockToolUse(BaseModel):
    """Tool use content block (model requesting a tool call)."""

    type: Literal["tool_use"] = "tool_use"
    id: str
    name: str
    input: dict[str, Any]


class ContentBlockToolResult(BaseModel):
    """Tool result content block (user providing tool output)."""

    type: Literal["tool_result"] = "tool_result"
    tool_use_id: str
    content: str | list[dict[str, Any]] | dict[str, Any] | list[Any] | Any
    is_error: bool | None = None
    cache_control: dict[str, str] | None = None


class ContentBlockThinking(BaseModel):
    """Thinking content block for reasoning models (e.g., Claude Opus 4.6)."""

    type: Literal["thinking"] = "thinking"
    thinking: str
    signature: str | None = None


class ContentBlockDocument(BaseModel):
    """Document content block (PDF, plain text)."""

    type: Literal["document"] = "document"
    source: dict[str, Any]  # {"type": "base64", "media_type": "application/pdf", "data": "..."}
    title: str | None = None
    context: str | None = None
    citations: dict[str, Any] | None = None
    cache_control: dict[str, str] | None = None


class ContentBlockInputAudio(BaseModel):
    """Audio input content block for multimodal audio models.

    Uses the same shape as OpenAI's input_audio content part so that
    the internal VLM engine receives a uniform representation regardless
    of which API endpoint the client used.
    """

    type: Literal["input_audio"] = "input_audio"
    input_audio: dict[str, Any]  # {"data": "<base64>", "format": "wav"}


# Union type for all content blocks
ContentBlock = (
    ContentBlockText
    | ContentBlockImage
    | ContentBlockToolUse
    | ContentBlockToolResult
    | ContentBlockThinking
    | ContentBlockDocument
    | ContentBlockInputAudio
)


# =============================================================================
# System Content
# =============================================================================


class SystemContent(BaseModel):
    """System message content block."""

    type: Literal["text"] = "text"
    text: str
    cache_control: dict[str, str] | None = None


# =============================================================================
# Messages
# =============================================================================


class AnthropicMessage(BaseModel):
    """A message in an Anthropic conversation."""

    role: Literal["user", "assistant", "system"]
    content: str | list[ContentBlock]


# =============================================================================
# Tools
# =============================================================================


class AnthropicTool(BaseModel):
    """Tool definition in Anthropic format.

    Supports two shapes:
      1. User-defined tool: requires ``input_schema``.
      2. Anthropic server-side tool (web_search, code_execution, bash,
         text_editor, computer): carries a versioned ``type`` like
         ``web_search_20250305`` and no ``input_schema``. oMLX cannot execute
         these locally; they are accepted for compatibility with clients such
         as Claude for Excel/PowerPoint/Word and dropped before inference.
    """

    name: str
    description: str | None = None
    input_schema: dict[str, Any] | None = None
    type: str | None = None
    cache_control: dict[str, str] | None = None

    # Forward-compat with extra fields Anthropic may attach to server-side
    # tools (e.g. max_uses, allowed_domains, user_location for web_search).
    model_config = ConfigDict(extra="allow")

    @model_validator(mode="after")
    def _require_schema_or_type(self) -> "AnthropicTool":
        if self.input_schema is None and self.type is None:
            raise ValueError(
                "AnthropicTool requires either 'input_schema' (user-defined "
                "tool) or 'type' (Anthropic server-side tool)."
            )
        return self


class ToolChoice(BaseModel):
    """Tool choice specification."""

    type: Literal["auto", "any", "tool"]
    name: str | None = None  # Required when type="tool"


# =============================================================================
# Thinking Configuration
# =============================================================================


class ThinkingConfig(BaseModel):
    """Configuration for extended thinking/reasoning."""

    type: Literal["enabled", "disabled", "adaptive"] = "enabled"
    budget_tokens: int | None = None


# =============================================================================
# Request
# =============================================================================


class MessagesRequest(BaseModel):
    """Request for Anthropic Messages API."""

    model: str
    max_tokens: int
    messages: list[AnthropicMessage]
    system: str | list[SystemContent] | None = None
    stop_sequences: list[str] | None = None
    stream: bool = False
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    metadata: dict[str, Any] | None = None
    tools: list[AnthropicTool] | None = None
    tool_choice: ToolChoice | dict[str, Any] | None = None
    thinking: ThinkingConfig | None = None
    # Chat template kwargs (e.g. enable_thinking, reasoning_effort)
    chat_template_kwargs: dict[str, Any] | None = None


# =============================================================================
# Token Counting
# =============================================================================


class TokenCountRequest(BaseModel):
    """Request for token counting (Anthropic format)."""

    model: str
    messages: list[AnthropicMessage]
    system: str | list[SystemContent] | None = None
    tools: list[AnthropicTool] | None = None
    tool_choice: ToolChoice | dict[str, Any] | None = None
    thinking: ThinkingConfig | None = None


class TokenCountResponse(BaseModel):
    """Response for token counting."""

    input_tokens: int


# =============================================================================
# Response
# =============================================================================


class AnthropicUsage(BaseModel):
    """Token usage statistics for Anthropic API."""

    input_tokens: int = 0
    output_tokens: int = 0
    cache_creation_input_tokens: int = 0
    cache_read_input_tokens: int = 0


class MessagesResponse(BaseModel):
    """Response for Anthropic Messages API."""

    id: str = Field(default_factory=lambda: generate_id(IDPrefix.MESSAGE))
    type: Literal["message"] = "message"
    role: Literal["assistant"] = "assistant"
    model: str
    content: list[ContentBlockText | ContentBlockToolUse | ContentBlockThinking]
    stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None = None
    stop_sequence: str | None = None
    usage: AnthropicUsage = Field(default_factory=AnthropicUsage)


# =============================================================================
# Streaming Events
# =============================================================================


class MessageStartEvent(BaseModel):
    """Event sent at the start of a message."""

    type: Literal["message_start"] = "message_start"
    message: dict[str, Any]  # Partial MessagesResponse


class ContentBlockStartEvent(BaseModel):
    """Event sent at the start of a content block."""

    type: Literal["content_block_start"] = "content_block_start"
    index: int
    content_block: dict[str, Any]  # Partial content block


class TextDelta(BaseModel):
    """Text delta for streaming."""

    type: Literal["text_delta"] = "text_delta"
    text: str


class InputJsonDelta(BaseModel):
    """JSON input delta for tool use streaming."""

    type: Literal["input_json_delta"] = "input_json_delta"
    partial_json: str


class ContentBlockDeltaEvent(BaseModel):
    """Event sent for content block updates."""

    type: Literal["content_block_delta"] = "content_block_delta"
    index: int
    delta: TextDelta | InputJsonDelta | dict[str, Any]


class ContentBlockStopEvent(BaseModel):
    """Event sent when a content block ends."""

    type: Literal["content_block_stop"] = "content_block_stop"
    index: int


class MessageDeltaEvent(BaseModel):
    """Event sent for message-level updates (stop_reason, usage)."""

    type: Literal["message_delta"] = "message_delta"
    delta: dict[str, Any]  # {"stop_reason": "...", "stop_sequence": ...}
    usage: dict[str, int]  # {"output_tokens": N}


class MessageStopEvent(BaseModel):
    """Event sent when the message ends."""

    type: Literal["message_stop"] = "message_stop"


class PingEvent(BaseModel):
    """Ping event for keeping connection alive."""

    type: Literal["ping"] = "ping"


class ErrorEvent(BaseModel):
    """Error event for streaming errors."""

    type: Literal["error"] = "error"
    error: dict[str, Any]  # {"type": "...", "message": "..."}


# Union type for all streaming events
StreamingEvent = (
    MessageStartEvent
    | ContentBlockStartEvent
    | ContentBlockDeltaEvent
    | ContentBlockStopEvent
    | MessageDeltaEvent
    | MessageStopEvent
    | PingEvent
    | ErrorEvent
)


# =============================================================================
# Error Response
# =============================================================================


class AnthropicErrorDetail(BaseModel):
    """Error detail in Anthropic format."""

    type: str  # "invalid_request_error", "authentication_error", "api_error", etc.
    message: str


class AnthropicErrorResponse(BaseModel):
    """Error response in Anthropic format."""

    type: Literal["error"] = "error"
    error: AnthropicErrorDetail


# ===== FILE: omlx/api/anthropic_utils.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Utility functions for Anthropic Messages API conversion.

Handles conversion between Anthropic API format and internal oMLX format.
"""

import base64
import json
import logging
import uuid
from typing import Any

from .anthropic_models import (
    AnthropicMessage,
    AnthropicTool,
    AnthropicUsage,
    ContentBlockText,
    ContentBlockThinking,
    ContentBlockToolUse,
    MessagesRequest,
    MessagesResponse,
    SystemContent,
)
from .openai_models import ToolCall

_PRESERVE_ROLE_BOUNDARY = "_preserve_role_boundary"
logger = logging.getLogger(__name__)


def request_has_cache_control(request: MessagesRequest) -> bool:
    """True if any system / tool / message block carries ``cache_control``.

    Anthropic's three input-side usage fields (``input_tokens``,
    ``cache_creation_input_tokens``, ``cache_read_input_tokens``) form a
    *disjoint* partition of the prompt only when the client explicitly
    marks a region with ``cache_control``. Without that signal the cache
    fields must stay at 0 and ``input_tokens`` carries the whole prompt
    count — independent of whether the oMLX engine happens to run
    automatic prefix caching internally.
    """
    sys = request.system
    if isinstance(sys, list):
        for blk in sys:
            if getattr(blk, "cache_control", None):
                return True

    for tool in request.tools or []:
        if getattr(tool, "cache_control", None):
            return True

    for msg in request.messages:
        content = msg.content
        if not isinstance(content, list):
            continue
        for blk in content:
            if getattr(blk, "cache_control", None):
                return True

    return False


def _decode_document_block(block_dict: dict[str, Any]) -> str:
    """Decode an Anthropic document content block to text.

    For text/plain documents, decodes base64 data and returns the text.
    For other media types (e.g. PDF), returns a placeholder message since
    oMLX does not provide document parsing.
    """
    source = block_dict.get("source", {})
    media_type = source.get("media_type", "")
    data = source.get("data", "")
    title = block_dict.get("title", "")

    if media_type == "text/plain" and data:
        try:
            decoded = base64.b64decode(data).decode("utf-8")
            label = f"[Document: {title}]\n" if title else ""
            return f"{label}{decoded}"
        except Exception:
            return f"[Document: {title or 'untitled'} — failed to decode]"

    label = title or "untitled"
    return (
        f"[Document: {label} ({media_type}) — "
        f"oMLX does not provide PDF parsing. Send as text instead.]"
    )


def _content_block_to_dict(block: Any) -> dict[str, Any] | None:
    """Normalize Anthropic content blocks to dicts."""
    if hasattr(block, "model_dump"):
        return block.model_dump()
    if isinstance(block, dict):
        return block
    return None


def _append_anthropic_image_part(
    image_parts: list[dict], block_dict: dict[str, Any]
) -> None:
    """Convert Anthropic image blocks to OpenAI-style image_url parts."""
    source = block_dict.get("source", {})
    if source.get("type") == "base64":
        media_type = source.get("media_type", "image/jpeg")
        data = source.get("data", "")
        image_parts.append(
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:{media_type};base64,{data}",
                },
            }
        )
    elif source.get("type") == "url":
        image_parts.append(
            {
                "type": "image_url",
                "image_url": {
                    "url": source.get("url", ""),
                },
            }
        )


def _append_anthropic_audio_part(
    audio_parts: list[dict], block_dict: dict[str, Any]
) -> None:
    """Pass through an input_audio block unchanged for the VLM engine."""
    input_audio = block_dict.get("input_audio")
    if input_audio and isinstance(input_audio, dict):
        audio_parts.append(
            {
                "type": "input_audio",
                "input_audio": input_audio,
            }
        )


def _extract_images_from_tool_result_content(
    content: Any, image_parts: list[dict]
) -> None:
    """Extract image blocks from tool result content for VLM processing."""
    if isinstance(content, list):
        for item in content:
            if isinstance(item, dict) and item.get("type") == "image":
                _append_anthropic_image_part(image_parts, item)
    elif isinstance(content, dict) and content.get("type") == "image":
        _append_anthropic_image_part(image_parts, content)


def _build_message_from_parts(
    role: str,
    text_parts: list[str],
    image_parts: list[dict],
    audio_parts: list[dict] | None = None,
) -> dict[str, Any] | None:
    """Build a single internal message from accumulated text/image/audio parts."""
    media_parts = list(image_parts)
    if audio_parts:
        media_parts.extend(audio_parts)

    if media_parts:
        content_parts = list(media_parts)
        if text_parts:
            content_parts.append(
                {
                    "type": "text",
                    "text": "\n".join(text_parts),
                }
            )
        return {"role": role, "content": content_parts}

    if text_parts:
        return {"role": role, "content": "\n".join(text_parts)}

    return None


# =============================================================================
# Message Conversion: Anthropic -> Internal
# =============================================================================


def convert_anthropic_to_internal(
    request: MessagesRequest,
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    preserve_images: bool = False,
    native_reasoning_content: bool = False,
    consolidate_system_messages: bool = True,
) -> list[dict[str, Any]]:
    """
    Convert Anthropic Messages API format to internal format.

    Handles:
    - System message from separate 'system' field
    - Content blocks to text
    - Tool results and tool uses in message history
    - Image blocks (when preserve_images=True for VLM)

    Args:
        request: Anthropic MessagesRequest
        max_tool_result_tokens: Maximum token count for tool results.
        tokenizer: Tokenizer instance for token counting and truncation.
        preserve_images: If True, preserve image blocks as OpenAI image_url
            format for VLM processing.
        native_reasoning_content: If True, attach Anthropic ``thinking`` blocks
            as a ``reasoning_content`` field on assistant messages (Qwen 3.6+
            templates).  If False, inline each block as ``<think>...</think>``
            in the message content as a fallback.
        consolidate_system_messages: If True, merge inline system messages into
            the leading system block. Server code can set this to False and let
            template capability probing decide whether mid-system messages can
            be preserved.

    Returns:
        List of {"role": str, "content": str or list}
    """
    from .utils import _chat_template_supports_tool_role

    processed_messages: list[dict[str, Any]] = []
    native_tool_calling = bool(
        tokenizer and _chat_template_supports_tool_role(tokenizer)
    )

    # Normalize: extract any role="system" entries from messages[] and merge
    # with the canonical request.system field (claude-code 2.1.154+ sends
    # system content inline instead of using the separate field).
    system_text, normalized_messages = _normalize_in_messages_system(
        request,
        consolidate_system_messages=consolidate_system_messages,
    )
    if system_text:
        processed_messages.append({"role": "system", "content": system_text})

    # Process messages
    for msg in normalized_messages:
        role = msg.role
        content = msg.content

        if isinstance(content, str):
            # Simple text message
            processed_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            if native_tool_calling:
                if role == "assistant":
                    text_parts: list[str] = []
                    image_parts: list[dict] = []
                    audio_parts: list[dict] = []
                    tool_calls: list[dict] = []
                    thinking_parts: list[str] = []
                    for block in content:
                        block_dict = _content_block_to_dict(block)
                        if block_dict is None:
                            continue
                        block_type = block_dict.get("type", "")
                        if block_type == "text":
                            text_parts.append(block_dict.get("text", ""))
                        elif block_type == "image" and preserve_images:
                            _append_anthropic_image_part(image_parts, block_dict)
                        elif block_type == "input_audio" and preserve_images:
                            _append_anthropic_audio_part(audio_parts, block_dict)
                        elif block_type == "tool_use":
                            tool_input = block_dict.get("input", {})
                            if isinstance(tool_input, str):
                                try:
                                    tool_input = json.loads(tool_input)
                                except (json.JSONDecodeError, ValueError):
                                    pass
                            tool_calls.append(
                                {
                                    "id": block_dict.get(
                                        "id", f"call_{uuid.uuid4().hex[:8]}"
                                    ),
                                    "function": {
                                        "name": block_dict.get("name", ""),
                                        "arguments": tool_input,
                                    },
                                }
                            )
                        elif block_type == "thinking":
                            # Native mode: collect for reasoning_content field.
                            # Fallback: inline as <think>...</think> in source
                            # order (Anthropic emits thinking first, so appending
                            # preserves the natural ordering).
                            thinking_text = block_dict.get("thinking", "")
                            if thinking_text:
                                if native_reasoning_content:
                                    thinking_parts.append(thinking_text)
                                else:
                                    text_parts.append(
                                        f"<think>\n{thinking_text}\n</think>"
                                    )
                        elif block_type == "document":
                            text_parts.append(_decode_document_block(block_dict))
                    msg_dict = _build_message_from_parts(
                        role, text_parts, image_parts, audio_parts
                    ) or {
                        "role": role,
                        "content": "",
                    }
                    if thinking_parts:
                        msg_dict["reasoning_content"] = "\n".join(thinking_parts)
                    if tool_calls:
                        msg_dict["tool_calls"] = tool_calls
                        msg_dict[_PRESERVE_ROLE_BOUNDARY] = True
                    processed_messages.append(msg_dict)
                    continue

                if role == "user":
                    text_parts = []
                    image_parts = []
                    audio_parts = []
                    saw_tool_result = False
                    for block in content:
                        block_dict = _content_block_to_dict(block)
                        if block_dict is None:
                            continue
                        block_type = block_dict.get("type", "")
                        if block_type == "text":
                            text_parts.append(block_dict.get("text", ""))
                        elif block_type == "image" and preserve_images:
                            _append_anthropic_image_part(image_parts, block_dict)
                        elif block_type == "input_audio" and preserve_images:
                            _append_anthropic_audio_part(audio_parts, block_dict)
                        elif block_type == "tool_result":
                            msg_dict = _build_message_from_parts(
                                role, text_parts, image_parts, audio_parts
                            )
                            if msg_dict:
                                processed_messages.append(msg_dict)
                            text_parts = []
                            image_parts = []
                            audio_parts = []
                            saw_tool_result = True
                            processed_messages.append(
                                {
                                    "role": "tool",
                                    "tool_call_id": block_dict.get("tool_use_id", ""),
                                    "content": _extract_tool_result_content(
                                        block_dict.get("content", ""),
                                        max_tokens=max_tool_result_tokens,
                                        tokenizer=tokenizer,
                                    ),
                                }
                            )
                            if preserve_images:
                                _extract_images_from_tool_result_content(
                                    block_dict.get("content", ""), image_parts
                                )
                        elif block_type == "thinking":
                            # User messages don't carry reasoning_content in the
                            # Qwen 3.6 template, so native mode simply drops these
                            # blocks.  Fallback keeps the legacy <think> inline
                            # behaviour in source order.
                            thinking_text = block_dict.get("thinking", "")
                            if thinking_text and not native_reasoning_content:
                                text_parts.append(f"<think>\n{thinking_text}\n</think>")
                        elif block_type == "document":
                            text_parts.append(_decode_document_block(block_dict))
                    msg_dict = _build_message_from_parts(
                        role, text_parts, image_parts, audio_parts
                    )
                    if msg_dict:
                        processed_messages.append(msg_dict)
                    elif not saw_tool_result:
                        processed_messages.append({"role": role, "content": ""})
                    continue

            # Content blocks list
            text_parts: list[str] = []
            image_parts: list[dict] = []
            audio_parts: list[dict] = []
            thinking_parts: list[str] = []
            saw_tool_markup = False
            for block in content:
                block_dict = _content_block_to_dict(block)
                if block_dict is None:
                    continue

                block_type = block_dict.get("type", "")

                if block_type == "text":
                    text_parts.append(block_dict.get("text", ""))

                elif block_type == "image" and preserve_images:
                    _append_anthropic_image_part(image_parts, block_dict)

                elif block_type == "input_audio" and preserve_images:
                    _append_anthropic_audio_part(audio_parts, block_dict)

                elif block_type == "tool_use":
                    # Tool use in assistant message (model called a tool)
                    tool_name = block_dict.get("name", "")
                    tool_input = block_dict.get("input", {})
                    text_parts.append(
                        f"[Calling tool: {tool_name}({json.dumps(tool_input)})]"
                    )
                    saw_tool_markup = True

                elif block_type == "tool_result":
                    # Tool result in user message (user providing tool output)
                    tool_use_id = block_dict.get("tool_use_id", "")
                    result_content = _extract_tool_result_content(
                        block_dict.get("content", ""),
                        max_tokens=max_tool_result_tokens,
                        tokenizer=tokenizer,
                    )
                    is_error = block_dict.get("is_error", False)
                    prefix = "[Tool Error" if is_error else "[Tool Result"
                    text_parts.append(f"{prefix} ({tool_use_id})]: {result_content}")
                    saw_tool_markup = True
                    if preserve_images:
                        _extract_images_from_tool_result_content(
                            block_dict.get("content", ""), image_parts
                        )

                elif block_type == "thinking":
                    # Native mode: collect for reasoning_content (assistant only).
                    # Fallback: inline as <think>...</think> in source order.
                    thinking_text = block_dict.get("thinking", "")
                    if thinking_text:
                        if native_reasoning_content and role == "assistant":
                            thinking_parts.append(thinking_text)
                        elif not native_reasoning_content:
                            text_parts.append(f"<think>\n{thinking_text}\n</think>")

                elif block_type == "document":
                    text_parts.append(_decode_document_block(block_dict))

            msg_dict = _build_message_from_parts(
                role, text_parts, image_parts, audio_parts
            ) or {
                "role": role,
                "content": "",
            }
            if thinking_parts:
                msg_dict["reasoning_content"] = "\n".join(thinking_parts)
            if saw_tool_markup:
                msg_dict[_PRESERVE_ROLE_BOUNDARY] = True
            processed_messages.append(msg_dict)
        else:
            # Unknown format
            processed_messages.append({"role": role, "content": str(content)})

    from .utils import _merge_consecutive_roles

    return _merge_consecutive_roles(processed_messages)


def convert_anthropic_to_internal_harmony(
    request: MessagesRequest,
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    consolidate_system_messages: bool = True,
) -> list[dict[str, Any]]:
    """
    Convert Anthropic Messages API format to internal format for Harmony (gpt-oss) models.

    Unlike convert_anthropic_to_internal(), this function preserves:
    - tool_use blocks as assistant.tool_calls field
    - tool_result blocks as role="tool" messages

    The Harmony chat_template expects these fields to properly generate
    the Harmony format tool calling syntax.

    Args:
        request: Anthropic MessagesRequest

    Returns:
        List of message dicts with tool-related fields preserved
    """
    processed_messages: list[dict[str, Any]] = []

    # Normalize: extract any role="system" entries from messages[] and merge
    # with the canonical request.system field (claude-code 2.1.154+ sends
    # system content inline instead of using the separate field).
    system_text, normalized_messages = _normalize_in_messages_system(
        request,
        consolidate_system_messages=consolidate_system_messages,
    )
    if system_text:
        processed_messages.append({"role": "system", "content": system_text})

    # Process messages
    for msg in normalized_messages:
        role = msg.role
        content = msg.content

        if isinstance(content, str):
            # Simple text message
            processed_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            # Content blocks list - need to separate tool_use, tool_result, and text
            text_parts: list[str] = []
            tool_calls: list[dict] = []
            tool_results: list[dict] = []

            for block in content:
                # Handle both Pydantic models and dicts
                if hasattr(block, "model_dump"):
                    block_dict = block.model_dump()
                elif isinstance(block, dict):
                    block_dict = block
                else:
                    continue

                block_type = block_dict.get("type", "")

                if block_type == "text":
                    text_parts.append(block_dict.get("text", ""))

                elif block_type == "tool_use":
                    # Tool use in assistant message - preserve as tool_calls
                    tool_id = block_dict.get("id", f"call_{uuid.uuid4().hex[:8]}")
                    tool_name = block_dict.get("name", "")
                    tool_input = block_dict.get("input", {})
                    # input should be dict for chat_template |tojson
                    if isinstance(tool_input, str):
                        try:
                            tool_input = json.loads(tool_input)
                        except (json.JSONDecodeError, ValueError):
                            pass
                    tool_calls.append(
                        {
                            "id": tool_id,
                            "function": {
                                "name": tool_name,
                                "arguments": tool_input,  # dict, not string
                            },
                        }
                    )

                elif block_type == "tool_result":
                    # Tool result - will be converted to role="tool" message
                    tool_use_id = block_dict.get("tool_use_id", "")
                    result_content = block_dict.get("content", "")

                    if isinstance(result_content, str):
                        # Try JSON parse BEFORE truncation so we can pretty-print
                        parsed_json = None
                        try:
                            parsed_json = json.loads(result_content)
                        except (json.JSONDecodeError, ValueError):
                            pass

                        if (
                            parsed_json is not None
                            and max_tool_result_tokens
                            and tokenizer
                        ):
                            # Valid JSON - pretty-print for better line-based truncation
                            pretty = json.dumps(
                                parsed_json, indent=2, ensure_ascii=False
                            )
                            truncated = truncate_tool_result(
                                pretty, max_tool_result_tokens, tokenizer
                            )
                            if "<truncated " in truncated:
                                # Truncation broke JSON - wrap in dict for
                                # Harmony |tojson compatibility
                                from .utils import _wrap_truncated_for_harmony

                                result_content = _wrap_truncated_for_harmony(truncated)
                            else:
                                # Not truncated - pass as parsed object
                                result_content = parsed_json
                        elif parsed_json is not None:
                            # Valid JSON, no truncation configured - pass as parsed object
                            result_content = parsed_json
                        else:
                            # Not JSON - apply truncation to raw text
                            result_content = _extract_tool_result_content(
                                result_content,
                                max_tokens=max_tool_result_tokens,
                                tokenizer=tokenizer,
                            )
                    elif isinstance(result_content, list):
                        # Extract text from content blocks
                        extracted = _extract_tool_result_content(
                            result_content,
                            max_tokens=max_tool_result_tokens,
                            tokenizer=tokenizer,
                        )
                        # Only try json.loads if content was NOT truncated
                        if (
                            isinstance(extracted, str)
                            and "<truncated " not in extracted
                        ):
                            try:
                                result_content = json.loads(extracted)
                            except (json.JSONDecodeError, ValueError):
                                result_content = extracted
                        elif isinstance(extracted, str) and "<truncated " in extracted:
                            # Check if pre-truncation content was JSON-like
                            content_part = extracted.split("\n\n<truncated")[0].strip()
                            if content_part and content_part[0] in "{[":
                                from .utils import _wrap_truncated_for_harmony

                                result_content = _wrap_truncated_for_harmony(extracted)
                            else:
                                result_content = extracted
                        else:
                            result_content = extracted
                    tool_results.append(
                        {
                            "tool_use_id": tool_use_id,
                            "content": result_content,
                        }
                    )

                elif block_type == "thinking":
                    # Thinking blocks are ignored (reasoning content is not passed to model)
                    continue

                elif block_type == "document":
                    text_parts.append(_decode_document_block(block_dict))

            # Build message(s) based on what we found
            if role == "assistant":
                # Assistant message with potential tool_calls
                msg_dict = {
                    "role": "assistant",
                    "content": "\n".join(text_parts) if text_parts else "",
                }
                if tool_calls:
                    msg_dict["tool_calls"] = tool_calls
                processed_messages.append(msg_dict)
            elif role == "user":
                # User message - may contain tool_results
                # First add any text content
                if text_parts:
                    processed_messages.append(
                        {"role": "user", "content": "\n".join(text_parts)}
                    )

                # Add each tool_result as a separate role="tool" message
                for tr in tool_results:
                    processed_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tr["tool_use_id"],
                            "content": tr["content"],  # dict or string
                        }
                    )
            else:
                # Other roles
                processed_messages.append(
                    {
                        "role": role,
                        "content": "\n".join(text_parts) if text_parts else "",
                    }
                )
        else:
            # Unknown format
            processed_messages.append({"role": role, "content": str(content)})

    from .utils import _merge_consecutive_roles

    return _merge_consecutive_roles(processed_messages)


# Prefix to filter out from system blocks (billing metadata that
# contains randomly changing values, breaking prefix cache).
_BILLING_HEADER_PREFIX = "x-anthropic-billing-header:"


def _extract_system_text(system: str | list[SystemContent]) -> str:
    """Extract text from system field."""
    if isinstance(system, str):
        return system
    elif isinstance(system, list):
        text_parts = []
        for block in system:
            if hasattr(block, "text"):
                text = block.text
            elif isinstance(block, dict) and block.get("type") == "text":
                text = block.get("text", "")
            else:
                continue
            # Skip billing header blocks (contain random values that break prefix cache)
            if text.startswith(_BILLING_HEADER_PREFIX):
                continue
            text_parts.append(text)
        return "\n".join(text_parts)
    return ""


def _normalize_in_messages_system(
    request: MessagesRequest,
    *,
    consolidate_system_messages: bool = True,
) -> tuple[str, list[AnthropicMessage]]:
    """Extract role="system" entries from messages[] and merge with request.system.

    Claude Code 2.1.154+ began sending system content inline in the messages
    array instead of (or in addition to) the canonical Anthropic ``system``
    field. Returns the combined system text and the message list with system
    entries removed, so downstream conversion sees the canonical shape.
    """
    if not consolidate_system_messages:
        base = _extract_system_text(request.system) if request.system else ""
        return base, list(request.messages)

    extracted_parts: list[str] = []
    filtered_messages: list[AnthropicMessage] = []
    for msg in request.messages:
        if msg.role != "system":
            filtered_messages.append(msg)
            continue
        content = msg.content
        if isinstance(content, str):
            if content:
                extracted_parts.append(content)
        elif isinstance(content, list):
            for block in content:
                block_dict = _content_block_to_dict(block)
                if block_dict is None:
                    continue
                if block_dict.get("type") == "text":
                    text = block_dict.get("text", "")
                    if text:
                        extracted_parts.append(text)

    base = _extract_system_text(request.system) if request.system else ""
    if extracted_parts:
        extra = "\n".join(extracted_parts)
        system_text = "\n\n".join(p for p in (base, extra) if p)
    else:
        system_text = base
    return system_text, filtered_messages


def truncate_tool_result(
    text: str,
    max_tokens: int,
    tokenizer: Any,
) -> str:
    """Truncate tool result text to fit within a token budget.

    Strategy:
    1. Encode the full text to count tokens.
    2. If within budget, return as-is.
    3. Decode tokens up to the budget to get an approximate character position.
    4. Search backwards for the last newline to truncate at a line boundary.
    5. Append a truncation notice as a separate XML tag.

    Args:
        text: The full tool result text.
        max_tokens: Maximum number of tokens allowed.
        tokenizer: Tokenizer with encode()/decode() methods.

    Returns:
        The (possibly truncated) text with notice appended.
    """
    token_ids = tokenizer.encode(text)
    total_tokens = len(token_ids)

    if total_tokens <= max_tokens:
        return text

    # Decode tokens up to budget to get approximate char position
    truncated_text = tokenizer.decode(token_ids[:max_tokens])

    # Find last newline for line-boundary truncation
    last_newline = truncated_text.rfind("\n")
    if last_newline > 0 and last_newline > len(truncated_text) * 0.5:
        # Only use line boundary if we don't lose more than 50% of content
        truncated_text = truncated_text[:last_newline]

    # Recount actual tokens after line-boundary adjustment
    shown_tokens = len(tokenizer.encode(truncated_text))

    logger.info(
        f"Tool result truncated: {total_tokens} -> {shown_tokens} tokens "
        f"({len(text)} -> {len(truncated_text)} chars)"
    )

    notice = (
        f'\n\n<truncated total_tokens="{total_tokens}" '
        f'shown_tokens="{shown_tokens}" />'
    )

    return truncated_text + notice


def _extract_tool_result_content(
    content: Any,
    max_tokens: int | None = None,
    tokenizer: Any | None = None,
) -> str:
    """Extract text from tool result content.

    Args:
        content: Raw tool result content (str, list, or dict).
        max_tokens: Maximum token count for the result. If exceeded, content is truncated.
        tokenizer: Tokenizer instance for token counting and truncation.

    Returns:
        Extracted text, potentially truncated if max_tokens is set.
    """
    if isinstance(content, str):
        result_text = content
    elif isinstance(content, list):
        # List of content blocks
        text_parts = []
        for item in content:
            if isinstance(item, dict):
                if item.get("type") == "text":
                    text_parts.append(item.get("text", ""))
            elif isinstance(item, str):
                text_parts.append(item)
        result_text = "\n".join(text_parts)
    elif isinstance(content, dict):
        if content.get("type") == "text":
            result_text = content.get("text", "")
        else:
            result_text = json.dumps(content)
    else:
        result_text = str(content)

    # Truncate by token count if configured
    if max_tokens and tokenizer and result_text:
        result_text = truncate_tool_result(result_text, max_tokens, tokenizer)
    elif max_tokens is not None:
        logger.debug(
            f"Tool result skip truncation: max_tokens={max_tokens}, "
            f"has_tokenizer={tokenizer is not None}, "
            f"result_len={len(result_text) if result_text else 0}"
        )

    return result_text


# =============================================================================
# Tool Conversion: Anthropic -> Internal
# =============================================================================

# Anthropic server-side tools (executed on Anthropic's infrastructure) carry a
# versioned ``type`` like ``web_search_20250305`` and have no ``input_schema``.
# oMLX cannot fulfill these locally, so we drop them before forwarding to the
# model. See https://docs.anthropic.com for the canonical tool families.
SERVER_SIDE_TOOL_TYPE_PREFIXES = (
    "web_search_",
    "code_execution_",
    "bash_",
    "text_editor_",
    "computer_",
)


def _is_server_side_tool(tool_dict: dict[str, Any]) -> bool:
    """Return True if the tool dict is an Anthropic server-side tool."""
    tool_type = tool_dict.get("type")
    if not isinstance(tool_type, str):
        return False
    return tool_type.startswith(SERVER_SIDE_TOOL_TYPE_PREFIXES)


def convert_anthropic_tools_to_internal(
    tools: list[AnthropicTool] | None,
) -> list[dict[str, Any]] | None:
    """
    Convert Anthropic tools to internal/OpenAI format.

    Anthropic: {"name": "...", "description": "...", "input_schema": {...}}
    Internal:  {"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}

    Anthropic server-side tools (web_search, code_execution, bash, text_editor,
    computer) cannot be executed by oMLX and are dropped with an INFO log.

    Args:
        tools: List of Anthropic tool definitions

    Returns:
        List of internal tool definitions, or None if no executable tools
    """
    if not tools:
        return None

    internal_tools: list[dict[str, Any]] = []
    dropped: list[str] = []

    for tool in tools:
        # Handle both Pydantic models and dicts
        if hasattr(tool, "model_dump"):
            tool_dict = tool.model_dump()
        elif isinstance(tool, dict):
            tool_dict = tool
        else:
            continue

        if _is_server_side_tool(tool_dict):
            dropped.append(f"{tool_dict.get('type')}:{tool_dict.get('name', '')}")
            continue

        internal_tools.append(
            {
                "type": "function",
                "function": {
                    "name": tool_dict.get("name", ""),
                    "description": tool_dict.get("description", ""),
                    "parameters": tool_dict.get("input_schema") or {},
                },
            }
        )

    if dropped:
        logger.info(
            "Dropped %d Anthropic server-side tool(s) not executable by oMLX: %s",
            len(dropped),
            ", ".join(dropped),
        )

    return internal_tools if internal_tools else None


# =============================================================================
# Response Conversion: Internal -> Anthropic
# =============================================================================


def convert_internal_to_anthropic_response(
    text: str,
    model: str,
    prompt_tokens: int,
    completion_tokens: int,
    finish_reason: str | None,
    tool_calls: list[ToolCall] | None = None,
    thinking: str | None = None,
    cached_tokens: int = 0,
    request_uses_cache_control: bool = False,
) -> MessagesResponse:
    """
    Convert internal output to Anthropic MessagesResponse.

    When the request carries ``cache_control`` breakpoints (signalled by
    ``request_uses_cache_control``) the prompt count is split into
    Anthropic's disjoint usage triple so that
    ``input_tokens + cache_creation_input_tokens + cache_read_input_tokens
    == prompt_tokens``. Otherwise the response keeps the legacy shape
    (``input_tokens = prompt_tokens``, both cache fields = 0) — even when
    the engine's automatic prefix cache happened to hit, since Anthropic
    only surfaces the cache triple when the client opted in.

    Args:
        text: Generated text content
        model: Model name
        prompt_tokens: Number of input tokens
        completion_tokens: Number of output tokens
        finish_reason: Internal finish reason ("stop", "length", "tool_calls")
        tool_calls: List of internal ToolCall objects
        thinking: Reasoning/thinking content from <think> blocks
        cached_tokens: Prompt tokens served from the prefix cache
        request_uses_cache_control: Whether the originating request carried
            ``cache_control`` on any system / tool / message block.

    Returns:
        Anthropic MessagesResponse
    """
    content: list[ContentBlockText | ContentBlockToolUse | ContentBlockThinking] = []

    # Add thinking content block before text if present.
    # Anthropic's spec requires a non-empty cryptographic signature on
    # thinking blocks; an empty string makes some SDK versions fall
    # back to a text-block parser path and emit "Content block is not
    # a text block". omlx cannot mint a real Anthropic signature, so
    # we use a stable placeholder string. Clients that strictly verify
    # the signature will still reject, but the common Claude Code SDK
    # only checks that the field is present and non-empty.
    if thinking and thinking.strip():
        content.append(
            ContentBlockThinking(
                type="thinking",
                thinking=thinking,
                signature="omlx-reasoning",
            )
        )

    # Add text content block if present and not empty
    if text and text.strip():
        content.append(ContentBlockText(type="text", text=text))

    # Add tool_use blocks if present
    if tool_calls:
        for tc in tool_calls:
            try:
                # Parse arguments from JSON string
                args = json.loads(tc.function.arguments)
            except (json.JSONDecodeError, AttributeError):
                args = {}

            content.append(
                ContentBlockToolUse(
                    type="tool_use",
                    id=tc.id,
                    name=tc.function.name,
                    input=args,
                )
            )

    # Ensure at least one content block
    if not content:
        content.append(ContentBlockText(type="text", text=""))

    # Map finish reason to stop reason
    stop_reason = map_finish_reason_to_stop_reason(finish_reason, bool(tool_calls))

    # Anthropic's three input-side fields are a disjoint partition of the
    # prompt and only carry non-zero values when the request opted into
    # caching via cache_control. Without that signal the cache fields stay
    # at 0 regardless of any internal prefix-cache hits in the engine.
    if request_uses_cache_control:
        cache_read = max(0, min(cached_tokens, prompt_tokens))
        cache_creation = prompt_tokens - cache_read
        input_display = 0
    else:
        cache_read = 0
        cache_creation = 0
        input_display = prompt_tokens

    return MessagesResponse(
        id=f"msg_{uuid.uuid4().hex[:24]}",
        type="message",
        role="assistant",
        model=model,
        content=content,
        stop_reason=stop_reason,
        usage=AnthropicUsage(
            input_tokens=input_display,
            output_tokens=completion_tokens,
            cache_creation_input_tokens=cache_creation,
            cache_read_input_tokens=cache_read,
        ),
    )


def map_finish_reason_to_stop_reason(
    finish_reason: str | None, has_tool_calls: bool
) -> str | None:
    """
    Map internal finish_reason to Anthropic stop_reason.

    Internal: "stop", "length", "tool_calls"
    Anthropic: "end_turn", "max_tokens", "stop_sequence", "tool_use"

    Args:
        finish_reason: Internal finish reason
        has_tool_calls: Whether the response contains tool calls

    Returns:
        Anthropic stop_reason
    """
    if has_tool_calls:
        return "tool_use"

    if finish_reason is None:
        return None

    mapping = {
        "stop": "end_turn",
        "length": "max_tokens",
        "tool_calls": "tool_use",
    }

    return mapping.get(finish_reason, "end_turn")


# =============================================================================
# SSE Event Formatting
# =============================================================================


def format_sse_event(event_type: str, data: dict[str, Any]) -> str:
    """
    Format an SSE event for Anthropic streaming.

    Anthropic uses: "event: {type}\\ndata: {json}\\n\\n"
    (Different from OpenAI which just uses "data: {json}\\n\\n")

    Args:
        event_type: Event type (message_start, content_block_delta, etc.)
        data: Event data to serialize as JSON

    Returns:
        Formatted SSE event string
    """
    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"


def create_message_start_event(
    message_id: str, model: str, input_tokens: int = 0
) -> str:
    """Create message_start SSE event."""
    return format_sse_event(
        "message_start",
        {
            "type": "message_start",
            "message": {
                "id": message_id,
                "type": "message",
                "role": "assistant",
                "model": model,
                "content": [],
                "stop_reason": None,
                "stop_sequence": None,
                "usage": {"input_tokens": input_tokens, "output_tokens": 0},
            },
        },
    )


def create_content_block_start_event(index: int, block_type: str, **kwargs) -> str:
    """Create content_block_start SSE event."""
    if block_type == "text":
        content_block = {"type": "text", "text": ""}
    elif block_type == "tool_use":
        content_block = {
            "type": "tool_use",
            "id": kwargs.get("id", ""),
            "name": kwargs.get("name", ""),
            "input": {},
        }
    elif block_type == "thinking":
        # Anthropic spec requires a signature field on thinking blocks
        # (see convert_internal_to_anthropic_response for the rationale
        # behind the placeholder string).
        content_block = {
            "type": "thinking",
            "thinking": "",
            "signature": "omlx-reasoning",
        }
    else:
        content_block = {"type": block_type}

    return format_sse_event(
        "content_block_start",
        {
            "type": "content_block_start",
            "index": index,
            "content_block": content_block,
        },
    )


def create_thinking_delta_event(index: int, thinking: str) -> str:
    """Create content_block_delta SSE event for thinking content."""
    return format_sse_event(
        "content_block_delta",
        {
            "type": "content_block_delta",
            "index": index,
            "delta": {"type": "thinking_delta", "thinking": thinking},
        },
    )


def create_text_delta_event(index: int, text: str) -> str:
    """Create content_block_delta SSE event for text."""
    return format_sse_event(
        "content_block_delta",
        {
            "type": "content_block_delta",
            "index": index,
            "delta": {"type": "text_delta", "text": text},
        },
    )


def create_input_json_delta_event(index: int, partial_json: str) -> str:
    """Create content_block_delta SSE event for tool input JSON."""
    return format_sse_event(
        "content_block_delta",
        {
            "type": "content_block_delta",
            "index": index,
            "delta": {"type": "input_json_delta", "partial_json": partial_json},
        },
    )


def create_content_block_stop_event(index: int) -> str:
    """Create content_block_stop SSE event."""
    return format_sse_event(
        "content_block_stop",
        {
            "type": "content_block_stop",
            "index": index,
        },
    )


def create_message_delta_event(
    stop_reason: str | None,
    output_tokens: int,
    stop_sequence: str | None = None,
    input_tokens: int | None = None,
    cached_tokens: int = 0,
    request_uses_cache_control: bool = False,
) -> str:
    """Create message_delta SSE event.

    When ``request_uses_cache_control`` is True and ``input_tokens`` is
    given, the count is split into Anthropic's disjoint triple (input
    stays 0, creation and read carry the remainder). Without that signal
    the cache fields are omitted entirely — Anthropic only surfaces them
    when the client opted in via a ``cache_control`` breakpoint, even if
    the engine's automatic prefix cache happened to hit.
    """
    usage: dict[str, int] = {"output_tokens": output_tokens}

    if request_uses_cache_control and input_tokens is not None:
        cache_read = max(0, min(cached_tokens, input_tokens))
        usage["input_tokens"] = 0
        usage["cache_creation_input_tokens"] = input_tokens - cache_read
        usage["cache_read_input_tokens"] = cache_read
    elif input_tokens is not None:
        usage["input_tokens"] = input_tokens

    return format_sse_event(
        "message_delta",
        {
            "type": "message_delta",
            "delta": {"stop_reason": stop_reason, "stop_sequence": stop_sequence},
            "usage": usage,
        },
    )


def create_message_stop_event() -> str:
    """Create message_stop SSE event."""
    return format_sse_event("message_stop", {"type": "message_stop"})


def create_ping_event() -> str:
    """Create ping SSE event."""
    return format_sse_event("ping", {"type": "ping"})


def create_error_event(error_type: str, message: str) -> str:
    """Create error SSE event."""
    return format_sse_event(
        "error",
        {
            "type": "error",
            "error": {"type": error_type, "message": message},
        },
    )


# ===== FILE: omlx/api/audio_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Pydantic models for OpenAI-compatible audio API.

These models define the request and response schemas for:
- Audio transcription (speech-to-text)
- Audio speech synthesis (text-to-speech)
"""

from typing import List, Optional

from pydantic import BaseModel


class AudioTranscriptionRequest(BaseModel):
    """OpenAI-compatible audio transcription request."""

    model: str
    language: Optional[str] = None
    prompt: Optional[str] = None
    response_format: Optional[str] = "json"
    temperature: Optional[float] = 0.0


class AudioTranscriptionResponse(BaseModel):
    text: str
    language: Optional[str] = None
    duration: Optional[float] = None
    segments: Optional[List[dict]] = None


class AudioSpeechRequest(BaseModel):
    model: str
    input: str
    voice: Optional[str] = None
    language: Optional[str] = None
    instructions: Optional[str] = None
    speed: Optional[float] = 1.0
    response_format: Optional[str] = "wav"
    ref_audio: Optional[str] = None
    ref_text: Optional[str] = None
    temperature: Optional[float] = None
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    repetition_penalty: Optional[float] = None
    max_tokens: Optional[int] = None
    stream: Optional[bool] = False
    streaming_interval: Optional[float] = None


class AudioProcessRequest(BaseModel):
    """Request model for audio processing (speech enhancement / STS).

    Used by POST /v1/audio/process — the audio file is submitted as a
    multipart upload alongside this model field.
    """

    model: str


# ===== FILE: omlx/api/audio_routes.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Audio API routes for oMLX.

This module provides OpenAI-compatible audio endpoints:
- POST /v1/audio/transcriptions  - Speech-to-Text
- POST /v1/audio/speech          - Text-to-Speech
- POST /v1/audio/process         - Speech-to-Speech / audio processing
"""

import asyncio
import base64
import io
import json
import logging
import math
import os
import re
import tempfile
from pathlib import Path
from typing import AsyncIterator, Optional

from fastapi import APIRouter, File, Form, HTTPException, UploadFile
from fastapi.responses import Response, StreamingResponse

from ..engine.audio_utils import wav_bytes_to_pcm_frames, wav_header
from ..server_metrics import get_server_metrics
from .audio_models import AudioSpeechRequest, AudioTranscriptionResponse

logger = logging.getLogger(__name__)

router = APIRouter()

# Maximum upload size for audio files (100 MB).
MAX_AUDIO_UPLOAD_BYTES = 100 * 1024 * 1024

# Maximum base64-encoded ref_audio size (~15 MB raw audio, enough for ~60s).
MAX_REF_AUDIO_BASE64_BYTES = 20 * 1024 * 1024

# Default native TTS chunk cadence. Keep this below the mlx-audio default to
# improve TTFT while still letting the model process the full input at once.
DEFAULT_NATIVE_TTS_STREAMING_INTERVAL_SECONDS = 0.2
MIN_NATIVE_TTS_STREAMING_INTERVAL_SECONDS = 0.01

# Non-streaming TTS output formats and their Content-Type. wav is the native
# engine output; the others are transcoded in memory by soundfile's bundled
# libsndfile (lame/opus/flac), which the audio extra already ships via
# mlx-audio -> librosa. aac is not offered (libsndfile has no aac encoder).
_SPEECH_RESPONSE_FORMATS = {
    "wav": "audio/wav",
    "mp3": "audio/mpeg",
    "opus": "audio/ogg",
    "flac": "audio/flac",
    "pcm": "audio/pcm",
}

# Video container extensions that should be routed through ffmpeg decoding.
# mlx-audio only recognises audio-specific extensions (m4a, aac, ogg, opus),
# so we remap video containers to .m4a before handing off. ffmpeg detects the
# actual format from file content, not the extension.
_VIDEO_CONTAINERS = {".mp4", ".mkv", ".mov", ".m4v", ".webm", ".avi"}


# ---------------------------------------------------------------------------
# Engine pool accessor — patched in tests via omlx.api.audio_routes._get_engine_pool
# ---------------------------------------------------------------------------


def _get_engine_pool():
    """Return the active EnginePool from server state.

    Imported lazily to avoid a circular import at module load time.
    Can be replaced in tests via patch('omlx.api.audio_routes._get_engine_pool').
    """
    # Import here to avoid circular imports at module load
    from omlx.server import _server_state

    pool = _server_state.engine_pool
    if pool is None:
        raise HTTPException(status_code=503, detail="Server not initialized")
    return pool


def _resolve_model(model_id: str) -> str:
    """Resolve a model alias to its real model ID.

    Delegates to the same resolve_model_id used by LLM/chat endpoints,
    ensuring audio endpoints handle aliases consistently.
    """
    from omlx.server import resolve_model_id

    return resolve_model_id(model_id) or model_id


def _get_settings_manager():
    """Return the active ModelSettingsManager from server state, or None.

    Lazy import + defensive guard so the audio router stays usable in tests
    that don't bring up the full server state.
    """
    try:
        from omlx.server import _server_state
    except Exception:
        return None
    return getattr(_server_state, "settings_manager", None)


def _record_audio_request(model_id: str) -> None:
    """Record audio request count without treating bytes/chars as tokens."""
    try:
        get_server_metrics().record_request_complete(
            prompt_tokens=0,
            completion_tokens=0,
            cached_tokens=0,
            model_id=model_id,
        )
    except Exception as exc:
        logger.warning("Failed to record audio metrics for %s: %s", model_id, exc)


async def _read_upload(file: UploadFile) -> bytes:
    """Read an uploaded file in chunks, bailing early if it exceeds the limit."""
    chunks: list[bytes] = []
    total = 0
    while True:
        chunk = await file.read(1024 * 1024)  # 1 MB chunks
        if not chunk:
            break
        total += len(chunk)
        if total > MAX_AUDIO_UPLOAD_BYTES:
            raise HTTPException(
                status_code=413,
                detail=(
                    f"Audio file exceeds maximum allowed size "
                    f"({MAX_AUDIO_UPLOAD_BYTES} bytes)"
                ),
            )
        chunks.append(chunk)
    return b"".join(chunks)


def _decode_ref_audio_base64(request: AudioSpeechRequest) -> Optional[bytes]:
    """Validate and decode optional base64 ref_audio from a TTS request."""
    if request.ref_audio is None:
        return None

    if not request.ref_text:
        raise HTTPException(
            status_code=400,
            detail="'ref_text' is required when 'ref_audio' is provided "
            "(must be the transcript of the reference audio)",
        )
    if len(request.ref_audio) > MAX_REF_AUDIO_BASE64_BYTES:
        raise HTTPException(
            status_code=413,
            detail=(
                f"ref_audio exceeds maximum allowed size "
                f"({MAX_REF_AUDIO_BASE64_BYTES} bytes base64, "
                f"~60 seconds of audio)"
            ),
        )
    try:
        return base64.b64decode(request.ref_audio, validate=True)
    except Exception:
        raise HTTPException(
            status_code=400,
            detail="Invalid base64 encoding in 'ref_audio' field",
        )


def _write_ref_audio_tempfile(audio_bytes: Optional[bytes]) -> Optional[str]:
    """Persist decoded ref audio to a temp file if present."""
    if audio_bytes is None:
        return None
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    try:
        tmp.write(audio_bytes)
        return tmp.name
    finally:
        tmp.close()


def _cleanup_tempfile(path: Optional[str]) -> None:
    if path and os.path.exists(path):
        try:
            os.unlink(path)
        except OSError:
            pass


def _resolve_tts_streaming_interval(request: AudioSpeechRequest) -> float:
    """Return a native TTS streaming interval that is safe for mlx-audio."""
    if request.streaming_interval is None:
        return DEFAULT_NATIVE_TTS_STREAMING_INTERVAL_SECONDS

    interval = request.streaming_interval
    if (
        not math.isfinite(interval)
        or interval < MIN_NATIVE_TTS_STREAMING_INTERVAL_SECONDS
    ):
        raise HTTPException(
            status_code=400,
            detail=(
                "'streaming_interval' must be at least "
                f"{MIN_NATIVE_TTS_STREAMING_INTERVAL_SECONDS} seconds"
            ),
        )
    return interval


def _split_tts_text(text: str, max_chars: int = 300) -> list[str]:
    """Split TTS input into conservative sentence-like chunks."""
    text = text.strip()
    if not text:
        return []

    sentences = re.split(r"(?<=[.!?。！？])\s+|\n+", text)
    sentences = [s.strip() for s in sentences if s and s.strip()]
    if not sentences:
        sentences = [text]

    chunks: list[str] = []
    current = ""

    def flush_current() -> None:
        nonlocal current
        if current:
            chunks.append(current.strip())
            current = ""

    for sentence in sentences:
        if len(sentence) > max_chars:
            flush_current()
            parts = re.split(r"(?<=[,;:，；：])\s*", sentence)
            parts = [p.strip() for p in parts if p and p.strip()]
            buffer = ""
            for part in parts or [sentence]:
                while len(part) > max_chars:
                    if buffer:
                        chunks.append(buffer.strip())
                        buffer = ""
                    chunks.append(part[:max_chars].strip())
                    part = part[max_chars:].strip()
                if not part:
                    continue
                candidate = f"{buffer} {part}".strip() if buffer else part
                if len(candidate) <= max_chars:
                    buffer = candidate
                else:
                    if buffer:
                        chunks.append(buffer.strip())
                    buffer = part
            if buffer:
                chunks.append(buffer.strip())
            continue

        candidate = f"{current} {sentence}".strip() if current else sentence
        if current and len(candidate) > max_chars:
            flush_current()
            current = sentence
        else:
            current = candidate

    flush_current()
    return chunks or [text]


async def _stream_speech_response(
    engine,
    request: AudioSpeechRequest,
    ref_audio_path: Optional[str],
    streaming_interval: float,
) -> AsyncIterator[bytes]:
    """Stream sentence-level TTS as a single WAV header plus PCM chunks."""
    try:
        if (
            hasattr(engine, "supports_native_tts_streaming")
            and engine.supports_native_tts_streaming()
            and hasattr(engine, "stream_synthesize_pcm")
        ):
            logger.info(
                "TTS native streaming start: model=%s, text_len=%d, voice=%s, language=%s",
                request.model, len(request.input), request.voice, request.language or "auto",
            )
            stream_format: Optional[tuple[int, int, int]] = None
            try:
                async for sample_rate, channels, sample_width, pcm_bytes in engine.stream_synthesize_pcm(
                    request.input,
                    voice=request.voice,
                    language=request.language,
                    speed=request.speed,
                    instructions=request.instructions,
                    ref_audio=ref_audio_path,
                    ref_text=request.ref_text,
                    temperature=request.temperature,
                    top_k=request.top_k,
                    top_p=request.top_p,
                    repetition_penalty=request.repetition_penalty,
                    max_tokens=request.max_tokens,
                    streaming_interval=streaming_interval,
                ):
                    fmt = (sample_rate, channels, sample_width)
                    if stream_format is None:
                        stream_format = fmt
                        yield wav_header(
                            sample_rate=sample_rate,
                            channels=channels,
                            sample_width=sample_width,
                        )
                    elif fmt != stream_format:
                        raise RuntimeError(
                            "Inconsistent native streaming PCM format: "
                            f"expected {stream_format}, got {fmt}"
                        )
                    if pcm_bytes:
                        yield pcm_bytes
            except NotImplementedError:
                if stream_format is not None:
                    raise
                logger.info(
                    "TTS native streaming unavailable at runtime; falling back "
                    "to segmented synthesis: model=%s",
                    request.model,
                )
            else:
                return

        segments = _split_tts_text(request.input)
        logger.info(
            "TTS streaming start: model=%s, text_len=%d, segments=%d, voice=%s, language=%s",
            request.model, len(request.input), len(segments), request.voice, request.language or "auto",
        )

        stream_format: Optional[tuple[int, int, int]] = None
        for idx, segment in enumerate(segments, start=1):
            wav_bytes = await engine.synthesize(
                segment,
                voice=request.voice,
                language=request.language,
                speed=request.speed,
                instructions=request.instructions,
                ref_audio=ref_audio_path,
                ref_text=request.ref_text,
                temperature=request.temperature,
                top_k=request.top_k,
                top_p=request.top_p,
                repetition_penalty=request.repetition_penalty,
                max_tokens=request.max_tokens,
            )
            sample_rate, channels, sample_width, pcm_bytes = wav_bytes_to_pcm_frames(wav_bytes)
            fmt = (sample_rate, channels, sample_width)
            if stream_format is None:
                stream_format = fmt
                yield wav_header(sample_rate=sample_rate, channels=channels, sample_width=sample_width)
            elif fmt != stream_format:
                raise RuntimeError(
                    "Inconsistent WAV format across TTS segments: "
                    f"expected {stream_format}, got {fmt}"
                )
            logger.debug(
                "TTS streaming segment %d/%d: text_len=%d, pcm_bytes=%d",
                idx, len(segments), len(segment), len(pcm_bytes),
            )
            if pcm_bytes:
                yield pcm_bytes
    finally:
        _cleanup_tempfile(ref_audio_path)


async def _stream_with_prefetched_chunk(
    first_chunk: bytes,
    stream: AsyncIterator[bytes],
) -> AsyncIterator[bytes]:
    """Yield a chunk fetched before response headers, then the rest of the stream."""
    try:
        yield first_chunk
        async for chunk in stream:
            yield chunk
    finally:
        close = getattr(stream, "aclose", None)
        if close is not None:
            await close()


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------


def _sse_event(payload: dict) -> str:
    """Serialize one data-only SSE event, OpenAI transcription-stream style."""
    return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"


async def _stream_transcription_events(
    engine,
    tmp_path: str,
    model_id: str,
    transcribe_kwargs: dict,
) -> AsyncIterator[str]:
    """Yield OpenAI-compatible transcript.text.* SSE events.

    Owns the uploaded temp file: it is deleted when the stream finishes,
    errors, or is cancelled by a client disconnect.
    """
    full_text: list[str] = []
    prompt_tokens = 0
    generation_tokens = 0
    try:
        async for chunk in engine.transcribe_stream(tmp_path, **transcribe_kwargs):
            # Cumulative totals arrive on the chunks that know them
            # (typically the final one); keep the max seen.
            prompt_tokens = max(
                prompt_tokens, int(chunk.get("prompt_tokens") or 0)
            )
            generation_tokens = max(
                generation_tokens, int(chunk.get("generation_tokens") or 0)
            )
            delta = chunk.get("text") or ""
            if not delta:
                continue
            full_text.append(delta)
            yield _sse_event({"type": "transcript.text.delta", "delta": delta})

        done: dict = {"type": "transcript.text.done", "text": "".join(full_text)}
        if prompt_tokens or generation_tokens:
            done["usage"] = {
                "type": "tokens",
                "input_tokens": prompt_tokens,
                "output_tokens": generation_tokens,
                "total_tokens": prompt_tokens + generation_tokens,
            }
        yield _sse_event(done)
        _record_audio_request(model_id)
    finally:
        _cleanup_tempfile(tmp_path)


@router.post("/v1/audio/transcriptions", response_model=AudioTranscriptionResponse)
async def create_transcription(
    file: UploadFile = File(...),
    model: str = Form(...),
    language: Optional[str] = Form(None),
    prompt: Optional[str] = Form(None),
    response_format: str = Form("json"),
    temperature: float = Form(0.0),
    stream: bool = Form(False),
    max_tokens: Optional[int] = Form(None),
    word_timestamps: bool = Form(False),
):
    """OpenAI-compatible audio transcription endpoint (Speech-to-Text).

    Note: ``response_format`` and ``temperature`` are accepted for OpenAI API
    compatibility but are not yet implemented — they are silently ignored.

    ``stream=true`` switches the response to OpenAI's transcription SSE
    format: ``transcript.text.delta`` events with incremental text followed
    by a final ``transcript.text.done`` event with the full transcription
    (#1066). Models whose mlx-audio backend lacks native streaming still
    respond in SSE format, with the full text arriving in a single delta.

    ``prompt`` follows the OpenAI spec: optional text to guide recognition
    toward specific vocabulary, spellings, or style. Mapped onto the
    backend's biasing hook — Qwen3-ASR receives it as trained context
    injection (``system_prompt``, strong biasing), Whisper models as a
    decoder-prefix soft prior (``initial_prompt``, ~224-token window).
    Backends without a biasing hook ignore it; it never fails a request.

    ``max_tokens`` is an oMLX extension that raises the underlying model's
    output cap. Useful for long audio with models like VibeVoice-ASR whose
    mlx-audio default (8192) truncates ~24 min files. When omitted, the
    model's own default applies.

    ``word_timestamps`` is an oMLX extension that exposes mlx-audio's native
    word-level alignment for Whisper models. When True, each segment in the
    response includes a ``words`` array of
    ``{word, start, end, probability}`` objects. Default False preserves the
    existing response shape for every current caller.
    """
    from omlx.engine.stt import STTEngine
    from omlx.exceptions import ModelNotFoundError

    pool = _get_engine_pool()
    resolved_model = _resolve_model(model)

    # Load the engine via pool (handles model loading and LRU eviction)
    try:
        engine = await pool.get_engine(resolved_model)
    except ModelNotFoundError as exc:
        avail = ", ".join(exc.available_models) if exc.available_models else "(none)"
        raise HTTPException(
            status_code=404,
            detail=f"Model '{resolved_model}' not found. Available: {avail}",
        ) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    if not isinstance(engine, STTEngine):
        raise HTTPException(
            status_code=400,
            detail=f"Model '{resolved_model}' is not a speech-to-text model",
        )

    # Save uploaded file to a temp path so the engine can open it by path.
    # Remap video container extensions to .m4a so mlx-audio routes them
    # through ffmpeg instead of miniaudio (which can't decode containers).
    suffix = os.path.splitext(file.filename or "audio.wav")[1] or ".wav"
    if suffix.lower() in _VIDEO_CONTAINERS:
        suffix = ".m4a"
    tmp_path = None
    try:
        content = await _read_upload(file)
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp_path = tmp.name
            tmp.write(content)

        # Effective max_tokens precedence: request > per-model setting (if any) >
        # model's own ``generate(max_tokens=...)`` default. The per-model lookup
        # mirrors how chat completions reads ModelSettings.max_tokens for LLMs;
        # for STT, settings.json's ``max_tokens`` (e.g. raised to 65536 for
        # VibeVoice-ASR) becomes the durable default for that model.
        effective_max_tokens = max_tokens
        if effective_max_tokens is None:
            sm = _get_settings_manager()
            if sm is not None:
                try:
                    ms = sm.get_settings(resolved_model)
                    if ms is not None and getattr(ms, "max_tokens", None) is not None:
                        effective_max_tokens = ms.max_tokens
                except Exception:
                    pass

        transcribe_kwargs: dict = {"language": language}
        if prompt is not None:
            transcribe_kwargs["prompt"] = prompt
        if effective_max_tokens is not None:
            transcribe_kwargs["max_tokens"] = effective_max_tokens
        if word_timestamps:
            transcribe_kwargs["word_timestamps"] = True

        if stream:
            # Word timestamps only exist in the JSON segment response;
            # SSE streaming emits plain text deltas (matching OpenAI, which
            # also rejects timestamp granularity with stream=true).
            transcribe_kwargs.pop("word_timestamps", None)
            # The event generator owns tmp_path from here: its finally block
            # deletes the file once the stream completes or errors — the
            # route's finally must not remove it while chunks are pending.
            events = _stream_transcription_events(
                engine, tmp_path, resolved_model, transcribe_kwargs
            )
            tmp_path = None
            first_event = await events.__anext__()
            return StreamingResponse(
                _stream_with_prefetched_chunk(first_event, events),
                media_type="text/event-stream",
                headers={"X-Accel-Buffering": "no", "Cache-Control": "no-cache"},
            )

        result = await engine.transcribe(tmp_path, **transcribe_kwargs)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

    _record_audio_request(resolved_model)

    # Build response directly from the dict returned by STTEngine
    segments = result.get("segments") or None

    return AudioTranscriptionResponse(
        text=result.get("text", ""),
        language=result.get("language"),
        duration=result.get("duration"),
        segments=segments,
    )


@router.get("/v1/audio/voices")
async def list_model_voices(model: Optional[str] = None):
    """List built-in speaker/voice names for a TTS model.

    Reads static metadata only — a ``voices/`` directory (Kokoro-style)
    or the speaker table in ``config.json`` (Qwen3-TTS CustomVoice's
    ``talker_config.spk_id``) — so the model does not need to be loaded.
    Returns ``{"model": ..., "voices": [...]}``; an empty list means the
    model has no named speakers (e.g. voice-cloning base models).
    """
    if not model:
        raise HTTPException(
            status_code=400, detail="'model' query parameter is required"
        )
    pool = _get_engine_pool()
    resolved = _resolve_model(model)
    entry = pool.get_entry(resolved)
    if entry is None:
        raise HTTPException(
            status_code=404, detail=f"Model '{resolved}' not found"
        )

    model_dir = Path(entry.model_path)
    voices: list[str] = []
    voices_dir = model_dir / "voices"
    if voices_dir.is_dir():
        voices = sorted({
            f.stem
            for f in voices_dir.iterdir()
            if f.suffix in (".safetensors", ".pt")
        })
    else:
        try:
            config = json.loads((model_dir / "config.json").read_text())
        except (OSError, ValueError):
            config = {}
        talker = config.get("talker_config") or {}
        spk = talker.get("spk_id") or config.get("spk_id") or {}
        if isinstance(spk, dict):
            voices = sorted(spk.keys())
    return {"model": resolved, "voices": voices}


def _transcode_speech_output(wav_bytes: bytes, response_format: str) -> bytes:
    """Transcode the engine's native WAV output to response_format in memory."""
    if response_format == "pcm":
        return wav_bytes_to_pcm_frames(wav_bytes)[3]

    # Lazy import like the other optional audio deps: soundfile ships with
    # the audio extra (mlx-audio -> librosa -> soundfile), and this module
    # must stay importable without it (see server.py route registration).
    import soundfile as sf

    data, sample_rate = sf.read(io.BytesIO(wav_bytes), dtype="float32")
    buf = io.BytesIO()
    if response_format == "mp3":
        sf.write(buf, data, sample_rate, format="MP3")
    elif response_format == "opus":
        sf.write(buf, data, sample_rate, format="OGG", subtype="OPUS")
    else:
        sf.write(buf, data, sample_rate, format="FLAC")
    return buf.getvalue()


@router.post("/v1/audio/speech")
async def create_speech(request: AudioSpeechRequest):
    """OpenAI-compatible text-to-speech endpoint."""
    from omlx.engine.tts import TTSEngine
    from omlx.exceptions import ModelNotFoundError

    if not request.input or not request.input.strip():
        raise HTTPException(status_code=400, detail="'input' field must not be empty")
    response_format = request.response_format or "wav"
    if response_format not in _SPEECH_RESPONSE_FORMATS:
        raise HTTPException(
            status_code=400,
            detail=(
                f"response_format '{request.response_format}' is not supported; "
                f"available formats: {', '.join(_SPEECH_RESPONSE_FORMATS)}"
            ),
        )
    streaming_interval = DEFAULT_NATIVE_TTS_STREAMING_INTERVAL_SECONDS
    if request.stream:
        if response_format != "wav":
            raise HTTPException(
                status_code=400,
                detail="Streaming TTS currently only supports response_format='wav'",
            )
        streaming_interval = _resolve_tts_streaming_interval(request)

    audio_bytes = _decode_ref_audio_base64(request)

    pool = _get_engine_pool()
    resolved_model = _resolve_model(request.model)

    try:
        engine = await pool.get_engine(resolved_model)
    except ModelNotFoundError as exc:
        avail = ", ".join(exc.available_models) if exc.available_models else "(none)"
        raise HTTPException(
            status_code=404,
            detail=f"Model '{resolved_model}' not found. Available: {avail}",
        ) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    if not isinstance(engine, TTSEngine):
        raise HTTPException(
            status_code=400,
            detail=f"Model '{resolved_model}' is not a text-to-speech model",
        )

    ref_audio_path = _write_ref_audio_tempfile(audio_bytes)

    if request.stream:
        stream = _stream_speech_response(
            engine,
            request,
            ref_audio_path,
            streaming_interval,
        )
        try:
            first_chunk = await stream.__anext__()
        except StopAsyncIteration as exc:
            raise HTTPException(
                status_code=500,
                detail="TTS streaming produced no audio output",
            ) from exc
        except HTTPException:
            raise
        except Exception as exc:
            raise HTTPException(status_code=500, detail=str(exc)) from exc
        return StreamingResponse(
            _stream_with_prefetched_chunk(first_chunk, stream),
            media_type="audio/wav",
        )

    try:
        wav_bytes = await engine.synthesize(
            request.input,
            voice=request.voice,
            language=request.language,
            speed=request.speed,
            instructions=request.instructions,
            ref_audio=ref_audio_path,
            ref_text=request.ref_text,
            temperature=request.temperature,
            top_k=request.top_k,
            top_p=request.top_p,
            repetition_penalty=request.repetition_penalty,
            max_tokens=request.max_tokens,
        )
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        _cleanup_tempfile(ref_audio_path)

    _record_audio_request(resolved_model)

    audio_out = wav_bytes
    if response_format != "wav":
        try:
            audio_out = await asyncio.to_thread(
                _transcode_speech_output, wav_bytes, response_format
            )
        except Exception as exc:
            raise HTTPException(status_code=500, detail=str(exc)) from exc

    return Response(
        content=audio_out, media_type=_SPEECH_RESPONSE_FORMATS[response_format]
    )


@router.post("/v1/audio/process")
async def process_audio(
    file: UploadFile = File(...),
    model: str = Form(...),
):
    """Audio processing endpoint (speech enhancement, source separation, STS).

    Accepts a multipart audio file upload and a model identifier, processes
    the audio through an STS engine (e.g. DeepFilterNet, MossFormer2,
    SAMAudio, LFM2.5-Audio), and returns WAV bytes of the processed audio.
    """
    from omlx.engine.sts import STSEngine
    from omlx.exceptions import ModelNotFoundError

    pool = _get_engine_pool()
    resolved_model = _resolve_model(model)

    # Load the engine via pool (handles model loading and LRU eviction)
    try:
        engine = await pool.get_engine(resolved_model)
    except ModelNotFoundError as exc:
        avail = ", ".join(exc.available_models) if exc.available_models else "(none)"
        raise HTTPException(
            status_code=404,
            detail=f"Model '{resolved_model}' not found. Available: {avail}",
        ) from exc
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    if not isinstance(engine, STSEngine):
        raise HTTPException(
            status_code=400,
            detail=f"Model '{resolved_model}' is not a speech-to-speech / audio processing model",
        )

    # Save uploaded file to a temp path so the engine can open it by path.
    # Remap video container extensions to .m4a so mlx-audio routes them
    # through ffmpeg instead of miniaudio (which can't decode containers).
    suffix = os.path.splitext(file.filename or "audio.wav")[1] or ".wav"
    if suffix.lower() in _VIDEO_CONTAINERS:
        suffix = ".m4a"
    tmp_path = None
    try:
        content = await _read_upload(file)
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp_path = tmp.name
            tmp.write(content)

        wav_bytes = await engine.process(tmp_path)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

    _record_audio_request(resolved_model)

    return Response(content=wav_bytes, media_type="audio/wav")


# ===== FILE: omlx/api/embedding_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Pydantic models for OpenAI-compatible Embeddings API.

These models define the request and response schemas for:
- /v1/embeddings endpoint
"""

import time
import uuid
from typing import List, Literal, Optional, Union

from pydantic import BaseModel, Field, model_validator


class EmbeddingInputItem(BaseModel):
    """Structured input item for multimodal embeddings."""

    text: Optional[str] = None
    # Image values are request-facing and must be inline data URIs. Remote URLs
    # and filesystem paths are rejected before processor-specific preparation.
    image: Optional[str] = None

    model_config = {"extra": "forbid"}

    @model_validator(mode="after")
    def validate_fields(self) -> "EmbeddingInputItem":
        """Require at least one supported field."""
        if self.text is None and self.image is None:
            raise ValueError("Embedding input item must include text or image")
        return self


class EmbeddingRequest(BaseModel):
    """
    Request for creating embeddings.

    OpenAI-compatible request format for the /v1/embeddings endpoint.
    """

    input: Optional[Union[str, List[str]]] = None
    """Input text(s) to embed. Can be a single string or list of strings."""

    items: Optional[List[EmbeddingInputItem]] = None
    """Structured embedding items for multimodal inputs."""

    model: str
    """ID of the model to use."""

    encoding_format: Literal["float", "base64"] = "float"
    """
    The format to return embeddings in.
    - "float": Returns a list of floats (default)
    - "base64": Returns a base64-encoded string of little-endian floats
    """

    dimensions: Optional[int] = None
    """
    The number of dimensions the output embeddings should have.
    Only supported by some models. If not supported, returns full dimensions.
    """

    max_length: Optional[int] = Field(default=None, gt=0)
    """
    Optional maximum token length for each input text. When omitted, the
    server uses the model's effective context window.
    """

    truncation: bool = True
    """
    Whether to truncate inputs longer than max_length.
    """

    @model_validator(mode="after")
    def validate_input_source(self) -> "EmbeddingRequest":
        """Require exactly one input source."""
        if self.input is None and self.items is None:
            raise ValueError("Either input or items must be provided")
        if self.input is not None and self.items is not None:
            raise ValueError("input and items cannot be provided together")
        if self.items is not None and len(self.items) == 0:
            raise ValueError("items cannot be empty")
        return self


class EmbeddingData(BaseModel):
    """A single embedding result."""

    object: str = "embedding"
    """The object type, always "embedding"."""

    index: int
    """The index of the embedding in the input list."""

    embedding: Union[List[float], str]
    """
    The embedding vector.
    - List[float] when encoding_format="float"
    - str (base64) when encoding_format="base64"
    """


class EmbeddingUsage(BaseModel):
    """Token usage statistics for embedding request."""

    prompt_tokens: int
    """Number of tokens in the input."""

    total_tokens: int
    """Total number of tokens used (same as prompt_tokens for embeddings)."""


class EmbeddingResponse(BaseModel):
    """
    Response from creating embeddings.

    OpenAI-compatible response format for the /v1/embeddings endpoint.
    """

    object: str = "list"
    """The object type, always "list"."""

    data: List[EmbeddingData]
    """List of embedding objects."""

    model: str
    """The model used for embedding."""

    usage: EmbeddingUsage
    """Usage statistics."""


# ===== FILE: omlx/api/embedding_utils.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Utility functions for the Embeddings API.

Provides:
- Base64 encoding for embeddings
- Dimension truncation with renormalization
- Token counting for usage statistics
"""

import base64
import math
import struct
from typing import Any, Dict, List, Union

from .embedding_models import EmbeddingInputItem


def encode_embedding_base64(embedding: List[float]) -> str:
    """
    Encode embedding vector as base64 string.

    OpenAI uses little-endian single-precision floats (float32).

    Args:
        embedding: List of float values

    Returns:
        Base64-encoded string of little-endian floats
    """
    packed = struct.pack(f"<{len(embedding)}f", *embedding)
    return base64.b64encode(packed).decode("ascii")


def truncate_embedding(embedding: List[float], dimensions: int) -> List[float]:
    """
    Truncate embedding to specified dimensions and renormalize.

    When truncating embeddings, we need to renormalize to maintain
    unit length (L2 norm = 1) for cosine similarity calculations.

    Args:
        embedding: Original embedding vector
        dimensions: Target number of dimensions

    Returns:
        Truncated and renormalized embedding
    """
    if dimensions >= len(embedding):
        return embedding

    truncated = embedding[:dimensions]

    # Calculate L2 norm
    norm = math.sqrt(sum(x * x for x in truncated))

    # Renormalize to unit length
    if norm > 0:
        return [x / norm for x in truncated]
    return truncated


def count_tokens(processor: Any, texts: List[str]) -> int:
    """
    Count total tokens in input texts.

    Handles different tokenizer/processor types from mlx-embeddings.

    Args:
        processor: Tokenizer or processor from mlx-embeddings
        texts: List of input texts

    Returns:
        Total number of tokens across all texts
    """
    total = 0

    for text in texts:
        # Try different encoding methods based on processor type
        if hasattr(processor, "encode"):
            # Standard tokenizer
            tokens = processor.encode(text, add_special_tokens=True)
            if isinstance(tokens, list):
                total += len(tokens)
            elif hasattr(tokens, "shape"):
                # MLX array
                total += tokens.shape[-1] if tokens.ndim > 0 else 1
            else:
                total += len(tokens)
        elif hasattr(processor, "tokenizer"):
            # Processor with nested tokenizer
            tokens = processor.tokenizer.encode(text, add_special_tokens=True)
            if isinstance(tokens, list):
                total += len(tokens)
            else:
                total += len(tokens)
        else:
            # Fallback: estimate based on whitespace
            total += len(text.split()) + 2  # +2 for special tokens

    return total


def normalize_input(input_data: Union[str, List[str]]) -> List[str]:
    """
    Normalize input to a list of strings.

    Args:
        input_data: Single string or list of strings

    Returns:
        List of strings
    """
    if isinstance(input_data, str):
        return [input_data]
    return list(input_data)


def normalize_embedding_items(
    items: List[Union[EmbeddingInputItem, Dict[str, Any]]]
) -> List[Dict[str, str]]:
    """
    Normalize structured embedding items into plain dicts.

    Args:
        items: Structured embedding input items

    Returns:
        List of normalized item dicts with only supported keys
    """
    normalized: List[Dict[str, str]] = []

    for item in items:
        if hasattr(item, "model_dump"):
            payload = item.model_dump(exclude_none=True)
        else:
            payload = {
                key: value for key, value in dict(item).items() if value is not None
            }

        text = payload.get("text")
        image = payload.get("image")

        normalized_item: Dict[str, str] = {}
        if text is not None:
            normalized_item["text"] = text
        if image is not None:
            normalized_item["image"] = image

        normalized.append(normalized_item)

    return normalized




# ===== FILE: omlx/api/grammar.py =====

# SPDX-License-Identifier: Apache-2.0
"""Grammar-constrained decoding via xgrammar.

Provides a logits processor that enforces grammar constraints by masking
invalid tokens at sampling time.  Follows the same ``__call__(tokens, logits)``
interface used by :class:`ThinkingBudgetProcessor`.

Phase-awareness (thinking vs. output) is handled by the *grammar itself*
via xgrammar's structural tag API, not by this processor.  For thinking
models the grammar is compiled as a ``sequence`` of
``[tag(<think>, any_text, </think>), constrained_schema]`` so that the
bitmask is permissive during reasoning and constrained during output.
This keeps the processor simple and enables uniform batched bitmask
computation (parallel model forward || bitmask fill).

The processor supports two usage modes:

1. **Per-request** (original): call ``processor(tokens, logits)`` directly.
   Handles accept + bitmask fill + mask application in one call.

2. **Batched**: call ``processor.advance(tokens)`` to accept the previous
   token, then use ``BatchGrammarMatcher.batch_fill_next_token_bitmask``
   with the exposed ``matcher`` property to fill bitmasks in parallel
   across the batch, and apply the combined bitmask externally.
"""

import logging
from typing import List, Optional

import mlx.core as mx
import numpy as np

logger = logging.getLogger(__name__)


def create_grammar_compiler(tokenizer, model):
    """Create an xgrammar GrammarCompiler for the given tokenizer and model.

    Returns None if vocab_size cannot be determined.
    """
    from .._torch_stub import install as _install_torch_stub
    _install_torch_stub()
    import xgrammar as xgr

    from ..utils.tokenizer import resolve_vocab_size, unwrap_tokenizer

    hf_tokenizer = unwrap_tokenizer(tokenizer)
    vocab_size = resolve_vocab_size(model)
    kwargs = {}
    if vocab_size is not None:
        kwargs["vocab_size"] = vocab_size

    tokenizer_info = xgr.TokenizerInfo.from_huggingface(hf_tokenizer, **kwargs)
    return xgr.GrammarCompiler(tokenizer_info)


class GrammarConstraintProcessor:
    """Logits processor that enforces grammar constraints via xgrammar bitmask.

    Args:
        compiled_grammar: An ``xgrammar.CompiledGrammar`` instance.  For
            thinking models this should already encode the thinking phase
            (compiled from a structural tag).
        vocab_size: Model vocabulary size (from model config, not tokenizer).
    """

    def __init__(self, compiled_grammar, vocab_size: int):
        from .._torch_stub import install as _install_torch_stub
        _install_torch_stub()
        import xgrammar as xgr
        from xgrammar.kernels.apply_token_bitmask_mlx import apply_token_bitmask_mlx

        self._matcher = xgr.GrammarMatcher(compiled_grammar)
        self._vocab_size = vocab_size
        self._apply_mask = apply_token_bitmask_mlx

        bitmask_width = (vocab_size + 31) // 32
        self._bitmask = np.full((1, bitmask_width), -1, dtype=np.int32)
        self._terminated = False
        self._first_call = True

    # ------------------------------------------------------------------
    # Per-request mode (original interface)
    # ------------------------------------------------------------------

    def __call__(self, tokens, logits: mx.array) -> mx.array:
        """Fill bitmask and apply to logits.

        Accept is handled by the monkey-patched GenerationBatch._step()
        which reads _next_tokens after sampling and calls accept_token().
        This method only fills the bitmask and applies it.
        """
        if self._terminated:
            return logits

        self._bitmask.fill(-1)
        self._matcher.fill_next_token_bitmask(self._bitmask)

        mx_bitmask = mx.array(self._bitmask)
        return self._apply_mask(mx_bitmask, logits, self._vocab_size)

    def accept_token(self, token_id: int) -> None:
        """Accept a generated token to advance matcher state."""
        if self._terminated:
            return
        if not self._matcher.accept_token(token_id):
            logger.warning("GrammarMatcher rejected token %d", token_id)
        if self._matcher.is_terminated():
            self._terminated = True

    # ------------------------------------------------------------------
    # Batched mode helpers
    # ------------------------------------------------------------------

    @property
    def matcher(self):
        """Return the underlying ``xgrammar.GrammarMatcher``."""
        return self._matcher

    @property
    def is_terminated(self) -> bool:
        return self._terminated

    def advance(self, tokens: mx.array) -> bool:
        """Accept the previous token and advance grammar state.

        Call this *instead of* ``__call__`` when using batched bitmask
        filling.  Returns ``True`` if the matcher is still active (not
        terminated) and should participate in the next
        ``batch_fill_next_token_bitmask`` call.
        """
        if self._terminated:
            return False

        if self._first_call:
            self._first_call = False
        elif len(tokens) > 0:
            last_token = int(tokens[-1])
            if not self._matcher.accept_token(last_token):
                logger.warning("GrammarMatcher rejected token %d", last_token)
            if self._matcher.is_terminated():
                self._terminated = True
                return False

        return True


# ===== FILE: omlx/api/markitdown.py =====

# SPDX-License-Identifier: Apache-2.0
"""MarkItDown integration helpers."""

from __future__ import annotations

import base64
import binascii
import io
import logging
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .openai_models import ContentPart, Message

logger = logging.getLogger(__name__)
NOISY_PDF_LOGGERS = (
    "pdfminer",
    "pdfminer.cmapdb",
    "pdfminer.converter",
    "pdfminer.layout",
    "pdfminer.pdfdocument",
    "pdfminer.pdfinterp",
    "pdfminer.pdfpage",
    "pdfminer.pdfparser",
    "pdfminer.psparser",
)

MARKITDOWN_MODEL_ID = "MarkItDown"
MARKITDOWN_MODEL_ALIASES = {"markitdown"}
MARKITDOWN_SUPPORTED_EXTENSIONS = {".pdf", ".docx", ".pptx"}
PLAIN_TEXT_SUPPORTED_EXTENSIONS = {".txt", ".md"}
SUPPORTED_ATTACHMENT_EXTENSIONS = (
    MARKITDOWN_SUPPORTED_EXTENSIONS | PLAIN_TEXT_SUPPORTED_EXTENSIONS
)
MARKITDOWN_UNSUPPORTED_SPREADSHEET_EXTENSIONS = {".xls", ".xlsx"}
DOCX_MIME_TYPE = (
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
)
PPTX_MIME_TYPE = (
    "application/vnd.openxmlformats-officedocument.presentationml.presentation"
)
MARKITDOWN_SUPPORTED_MIME_TYPES = {
    "application/pdf",
    DOCX_MIME_TYPE,
    PPTX_MIME_TYPE,
}
PLAIN_TEXT_SUPPORTED_MIME_TYPES = {"text/plain", "text/markdown", "text/x-markdown"}
SUPPORTED_ATTACHMENT_MIME_TYPES = (
    MARKITDOWN_SUPPORTED_MIME_TYPES | PLAIN_TEXT_SUPPORTED_MIME_TYPES
)
MARKITDOWN_SPREADSHEET_MIME_TYPES = {
    "application/vnd.ms-excel",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
}
MARKITDOWN_EMPTY_PDF_MESSAGE = (
    "No extractable text found. This PDF appears to be scanned or image-only. "
    "Select an OCR PDF processing engine for scanned or image-only PDFs."
)
MARKITDOWN_EMPTY_MESSAGE = "No extractable text found in the attached document."
PDF_PROCESSING_MARKITDOWN = "markitdown"
DEFAULT_PDF_PROCESSING_ENGINE = PDF_PROCESSING_MARKITDOWN

_converter_lock = threading.Lock()
_converter: Any | None = None


@dataclass(frozen=True)
class MarkItDownFile:
    filename: str
    mime_type: str
    data: bytes

    @property
    def extension(self) -> str:
        return Path(self.filename).suffix.lower()


class MarkItDownRequestError(ValueError):
    """Client-side request error for document conversion."""

    def __init__(self, detail: str, *, status_code: int = 400) -> None:
        super().__init__(detail)
        self.detail = detail
        self.status_code = status_code


def is_markitdown_model(model: str | None) -> bool:
    normalized = (model or "").strip()
    return normalized == MARKITDOWN_MODEL_ID or normalized in MARKITDOWN_MODEL_ALIASES


def markitdown_enabled(global_settings: Any | None) -> bool:
    if global_settings is None:
        return True
    integrations = getattr(global_settings, "integrations", None)
    if integrations is None:
        return True
    return bool(getattr(integrations, "markitdown_enabled", True))


def markitdown_model_visible(global_settings: Any | None) -> bool:
    if not markitdown_enabled(global_settings):
        return False
    integrations = getattr(global_settings, "integrations", None)
    if integrations is None:
        return False
    return bool(getattr(integrations, "markitdown_expose_model", False))


def markitdown_limits(global_settings: Any | None) -> tuple[int, int]:
    integrations = getattr(global_settings, "integrations", None)
    max_mb = getattr(integrations, "markitdown_max_file_size_mb", 25)
    max_files = getattr(integrations, "markitdown_max_files_per_request", 5)
    return int(max_mb or 25), int(max_files or 5)


def quiet_pdf_parser_loggers() -> None:
    for logger_name in NOISY_PDF_LOGGERS:
        logging.getLogger(logger_name).setLevel(logging.WARNING)


def markitdown_pdf_processing_engine(global_settings: Any | None) -> str:
    integrations = getattr(global_settings, "integrations", None)
    value = getattr(
        integrations,
        "markitdown_pdf_processing_engine",
        DEFAULT_PDF_PROCESSING_ENGINE,
    )
    return normalize_pdf_processing_engine(value)


def normalize_pdf_processing_engine(value: Any) -> str:
    engine = str(value or DEFAULT_PDF_PROCESSING_ENGINE).strip()
    if not engine:
        return DEFAULT_PDF_PROCESSING_ENGINE

    if engine.lower() == PDF_PROCESSING_MARKITDOWN:
        return PDF_PROCESSING_MARKITDOWN
    return engine


def request_has_file_parts(messages: list[Message]) -> bool:
    return any(_iter_file_part_dicts(msg.content) for msg in messages)


def convert_messages_to_markdown(
    messages: list[Message],
    *,
    global_settings: Any | None = None,
    latest_user_only: bool = False,
) -> str:
    """Render request messages into Markdown, converting file parts."""
    if latest_user_only:
        messages = _latest_user_turn(messages)

    converted_messages = preprocess_markitdown_file_parts(
        messages,
        global_settings=global_settings,
        fail_when_disabled=True,
    )
    sections: list[str] = []
    for msg in converted_messages:
        text = _content_text(msg.content).strip()
        if text:
            sections.append(text)
    return "\n\n".join(sections).strip()


async def convert_messages_to_markdown_async(
    messages: list[Message],
    *,
    global_settings: Any | None = None,
    engine_pool: Any | None = None,
    settings_manager: Any | None = None,
    get_sampling_params: Any | None = None,
    latest_user_only: bool = False,
) -> str:
    """Render request messages into Markdown, converting file parts asynchronously."""
    if latest_user_only:
        messages = _latest_user_turn(messages)

    converted_messages = await preprocess_markitdown_file_parts_async(
        messages,
        global_settings=global_settings,
        engine_pool=engine_pool,
        settings_manager=settings_manager,
        get_sampling_params=get_sampling_params,
        fail_when_disabled=True,
    )
    sections: list[str] = []
    for msg in converted_messages:
        text = _content_text(msg.content).strip()
        if text:
            sections.append(text)
    return "\n\n".join(sections).strip()


async def stream_messages_to_markdown_async(
    messages: list[Message],
    *,
    global_settings: Any | None = None,
    engine_pool: Any | None = None,
    settings_manager: Any | None = None,
    get_sampling_params: Any | None = None,
    latest_user_only: bool = False,
):
    """Stream request messages into Markdown chunks, preserving OCR PDF page order."""
    if latest_user_only:
        messages = _latest_user_turn(messages)

    if request_has_file_parts(messages) and not markitdown_enabled(global_settings):
        raise MarkItDownRequestError(
            "MarkItDown integration is disabled.",
            status_code=400,
        )

    max_file_size_mb, max_files = markitdown_limits(global_settings)
    files_seen = 0
    emitted = False

    async def emit_chunk_stream(chunks):
        nonlocal emitted
        started = False
        async for chunk in chunks:
            if not chunk:
                continue
            if not started:
                if emitted:
                    yield "\n\n"
                emitted = True
                started = True
            yield chunk

    for msg in messages:
        content = msg.content
        if isinstance(content, str):
            async for chunk in emit_chunk_stream(_single_chunk(content.strip())):
                yield chunk
            continue
        if not isinstance(content, list):
            async for chunk in emit_chunk_stream(_single_chunk(str(content))):
                yield chunk
            continue

        for part in content:
            part_dict = _part_to_dict(part)
            if part_dict.get("type") in {"text", "input_text"}:
                async for chunk in emit_chunk_stream(
                    _single_chunk(str(part_dict.get("text") or "").strip())
                ):
                    yield chunk
                continue

            if part_dict.get("type") != "file":
                continue

            files_seen += 1
            if files_seen > max_files:
                raise MarkItDownRequestError(
                    f"Too many attached files. Maximum is {max_files} per request.",
                    status_code=400,
                )

            parsed = parse_file_part(part_dict, max_file_size_mb=max_file_size_mb)
            async for chunk in emit_chunk_stream(
                stream_attachment_to_markdown_async(
                    parsed,
                    global_settings=global_settings,
                    engine_pool=engine_pool,
                    settings_manager=settings_manager,
                    get_sampling_params=get_sampling_params,
                )
            ):
                yield chunk


async def _single_chunk(text: str):
    if text:
        yield text


def _latest_user_turn(messages: list[Message]) -> list[Message]:
    for msg in reversed(messages):
        if msg.role == "user":
            return [msg]
    return messages[-1:] if messages else []


def preprocess_markitdown_file_parts(
    messages: list[Message],
    *,
    global_settings: Any | None = None,
    fail_when_disabled: bool = True,
    allow_missing_historical_files: bool = False,
) -> list[Message]:
    """Replace file content parts with Markdown text parts."""
    if not request_has_file_parts(messages):
        return messages

    if not markitdown_enabled(global_settings):
        if fail_when_disabled:
            raise MarkItDownRequestError(
                "MarkItDown integration is disabled.",
                status_code=400,
            )
        return messages

    max_file_size_mb, max_files = markitdown_limits(global_settings)
    files_seen = 0
    processed: list[Message] = []

    latest_user_index = _latest_user_index(messages)

    for msg_index, msg in enumerate(messages):
        content = msg.content
        if not isinstance(content, list):
            processed.append(msg)
            continue

        new_parts: list[ContentPart] = []
        changed = False
        for part in content:
            part_dict = _part_to_dict(part)
            if part_dict.get("type") != "file":
                new_parts.append(ContentPart.model_validate(part_dict))
                continue

            if allow_missing_historical_files and _is_missing_historical_file_part(
                part_dict,
                msg=msg,
                msg_index=msg_index,
                latest_user_index=latest_user_index,
            ):
                new_parts.append(
                    ContentPart(
                        type="text",
                        text=_format_missing_attachment(part_dict),
                    )
                )
                changed = True
                continue

            files_seen += 1
            if files_seen > max_files:
                raise MarkItDownRequestError(
                    f"Too many attached files. Maximum is {max_files} per request.",
                    status_code=400,
                )

            parsed = parse_file_part(part_dict, max_file_size_mb=max_file_size_mb)
            markdown = convert_attachment_to_markdown(
                parsed,
                global_settings=global_settings,
            )
            new_parts.append(
                ContentPart(
                    type="text",
                    text=_format_attachment_markdown(parsed.filename, markdown),
                )
            )
            changed = True

        processed.append(
            msg.model_copy(update={"content": new_parts}) if changed else msg
        )

    return processed


async def preprocess_markitdown_file_parts_async(
    messages: list[Message],
    *,
    global_settings: Any | None = None,
    engine_pool: Any | None = None,
    settings_manager: Any | None = None,
    get_sampling_params: Any | None = None,
    fail_when_disabled: bool = True,
    allow_missing_historical_files: bool = False,
) -> list[Message]:
    """Replace file content parts with Markdown text parts asynchronously."""
    if not request_has_file_parts(messages):
        return messages

    if not markitdown_enabled(global_settings):
        if fail_when_disabled:
            raise MarkItDownRequestError(
                "MarkItDown integration is disabled.",
                status_code=400,
            )
        return messages

    max_file_size_mb, max_files = markitdown_limits(global_settings)
    files_seen = 0
    processed: list[Message] = []

    latest_user_index = _latest_user_index(messages)

    for msg_index, msg in enumerate(messages):
        content = msg.content
        if not isinstance(content, list):
            processed.append(msg)
            continue

        new_parts: list[ContentPart] = []
        changed = False
        for part in content:
            part_dict = _part_to_dict(part)
            if part_dict.get("type") != "file":
                new_parts.append(ContentPart.model_validate(part_dict))
                continue

            if allow_missing_historical_files and _is_missing_historical_file_part(
                part_dict,
                msg=msg,
                msg_index=msg_index,
                latest_user_index=latest_user_index,
            ):
                new_parts.append(
                    ContentPart(
                        type="text",
                        text=_format_missing_attachment(part_dict),
                    )
                )
                changed = True
                continue

            files_seen += 1
            if files_seen > max_files:
                raise MarkItDownRequestError(
                    f"Too many attached files. Maximum is {max_files} per request.",
                    status_code=400,
                )

            parsed = parse_file_part(part_dict, max_file_size_mb=max_file_size_mb)
            markdown = await convert_attachment_to_markdown_async(
                parsed,
                global_settings=global_settings,
                engine_pool=engine_pool,
                settings_manager=settings_manager,
                get_sampling_params=get_sampling_params,
            )
            new_parts.append(
                ContentPart(
                    type="text",
                    text=_format_attachment_markdown(parsed.filename, markdown),
                )
            )
            changed = True

        processed.append(
            msg.model_copy(update={"content": new_parts}) if changed else msg
        )

    return processed


def parse_file_part(part: dict[str, Any], *, max_file_size_mb: int) -> MarkItDownFile:
    file_obj = part.get("file")
    if not isinstance(file_obj, dict):
        raise MarkItDownRequestError("File content part must include a file object.")

    data_value = file_obj.get("file_data") or file_obj.get("data")
    if not data_value and file_obj.get("file_id"):
        raise MarkItDownRequestError(
            "File content part file_id is not supported. "
            "Send base64 content in file.file_data instead.",
            status_code=400,
        )
    if not isinstance(data_value, str) or not data_value.strip():
        raise MarkItDownRequestError(
            "File content part requires file.file_data.",
            status_code=400,
        )

    data_uri_mime_type = _mime_type_from_data_uri(data_value)
    mime_type = (
        str(file_obj.get("mime_type") or data_uri_mime_type or "").strip().lower()
    )

    filename = str(file_obj.get("filename") or "").strip()
    if not filename:
        extension = _extension_for_mime_type(mime_type)
        if not extension:
            raise MarkItDownRequestError(
                "File content part requires file.filename.",
                status_code=400,
            )
        filename = f"attachment{extension}"

    extension = Path(filename).suffix.lower()
    _validate_supported_file(filename, mime_type)

    data = _decode_data(data_value)
    max_bytes = max_file_size_mb * 1024 * 1024
    if len(data) > max_bytes:
        raise MarkItDownRequestError(
            f"Attached file exceeds the {max_file_size_mb}MB limit: {filename}",
            status_code=400,
        )

    if not mime_type:
        mime_type = _mime_type_for_extension(extension)

    return MarkItDownFile(filename=filename, mime_type=mime_type, data=data)


def _latest_user_index(messages: list[Message]) -> int:
    for index in range(len(messages) - 1, -1, -1):
        if messages[index].role == "user":
            return index
    return -1


def _is_missing_historical_file_part(
    part: dict[str, Any],
    *,
    msg: Message,
    msg_index: int,
    latest_user_index: int,
) -> bool:
    if msg.role != "user" or msg_index >= latest_user_index:
        return False
    file_obj = part.get("file")
    if not isinstance(file_obj, dict):
        return False
    data_value = file_obj.get("file_data") or file_obj.get("data")
    return not isinstance(data_value, str) or not data_value.strip()


def _format_missing_attachment(part: dict[str, Any]) -> str:
    file_obj = part.get("file") if isinstance(part.get("file"), dict) else {}
    filename = str(file_obj.get("filename") or "attachment").strip()
    return f"## Attached file unavailable: {filename}"


def convert_attachment_to_markdown(
    file: MarkItDownFile,
    *,
    global_settings: Any | None = None,
) -> str:
    if _is_plain_text_attachment(file):
        return decode_plain_text_attachment(file)
    return convert_file_to_markdown(file, global_settings=global_settings)


async def convert_attachment_to_markdown_async(
    file: MarkItDownFile,
    *,
    global_settings: Any | None = None,
    engine_pool: Any | None = None,
    settings_manager: Any | None = None,
    get_sampling_params: Any | None = None,
) -> str:
    if _is_plain_text_attachment(file):
        return decode_plain_text_attachment(file)

    if file.extension == ".pdf":
        pdf_engine = markitdown_pdf_processing_engine(global_settings)
        if pdf_engine != PDF_PROCESSING_MARKITDOWN:
            from .markitdown_pdf_fallback import convert_pdf_with_ocr_engine

            return await convert_pdf_with_ocr_engine(
                file,
                engine_model_id=pdf_engine,
                engine_pool=engine_pool,
                settings_manager=settings_manager,
                global_settings=global_settings,
                get_sampling_params=get_sampling_params,
            )

    import asyncio

    return await asyncio.to_thread(
        convert_file_to_markdown,
        file,
        global_settings=global_settings,
    )


async def stream_attachment_to_markdown_async(
    file: MarkItDownFile,
    *,
    global_settings: Any | None = None,
    engine_pool: Any | None = None,
    settings_manager: Any | None = None,
    get_sampling_params: Any | None = None,
):
    if _is_plain_text_attachment(file):
        markdown = decode_plain_text_attachment(file)
        yield _format_attachment_markdown(file.filename, markdown)
        return

    if file.extension == ".pdf":
        pdf_engine = markitdown_pdf_processing_engine(global_settings)
        if pdf_engine != PDF_PROCESSING_MARKITDOWN:
            from .markitdown_pdf_fallback import stream_pdf_with_ocr_engine

            yield f"## Attached file: {file.filename}\n\n"
            async for chunk in stream_pdf_with_ocr_engine(
                file,
                engine_model_id=pdf_engine,
                engine_pool=engine_pool,
                settings_manager=settings_manager,
                global_settings=global_settings,
                get_sampling_params=get_sampling_params,
            ):
                yield chunk
            return

    markdown = await convert_attachment_to_markdown_async(
        file,
        global_settings=global_settings,
        engine_pool=engine_pool,
        settings_manager=settings_manager,
        get_sampling_params=get_sampling_params,
    )
    yield _format_attachment_markdown(file.filename, markdown)


def decode_plain_text_attachment(file: MarkItDownFile) -> str:
    try:
        text = file.data.decode("utf-8-sig")
    except UnicodeDecodeError as exc:
        raise MarkItDownRequestError(
            f"Text attachment must be UTF-8 encoded: {file.filename}",
            status_code=400,
        ) from exc

    markdown = text.strip()
    if not markdown:
        raise MarkItDownRequestError(MARKITDOWN_EMPTY_MESSAGE, status_code=400)
    return markdown


def convert_file_to_markdown(
    file: MarkItDownFile,
    *,
    global_settings: Any | None = None,
) -> str:
    if file.extension == ".pdf":
        return convert_pdf_file_to_markdown(file)

    markdown = _convert_file_with_markitdown(file)
    if not _has_extractable_text(markdown):
        logger.warning(
            "MarkItDown found no extractable text: filename=%s mime_type=%s",
            file.filename,
            file.mime_type,
        )
        raise MarkItDownRequestError(MARKITDOWN_EMPTY_MESSAGE, status_code=400)

    return markdown


def convert_pdf_file_to_markdown(
    file: MarkItDownFile,
) -> str:
    quiet_pdf_parser_loggers()
    markdown = _convert_file_with_markitdown(file)
    if not _has_extractable_text(markdown):
        logger.warning(
            "MarkItDown found no extractable text: filename=%s mime_type=%s",
            file.filename,
            file.mime_type,
        )
        raise MarkItDownRequestError(MARKITDOWN_EMPTY_PDF_MESSAGE, status_code=400)
    return markdown


def _convert_file_with_markitdown(file: MarkItDownFile) -> str:
    converter = _get_converter()

    try:
        from markitdown import StreamInfo
    except ImportError as exc:
        raise RuntimeError(
            "MarkItDown is not installed. Install markitdown[pdf,docx,pptx]."
        ) from exc

    stream = io.BytesIO(file.data)
    stream_info = StreamInfo(
        extension=file.extension,
        mimetype=file.mime_type or None,
        filename=file.filename,
    )

    try:
        with _converter_lock:
            result = converter.convert_stream(stream, stream_info=stream_info)
    except Exception as exc:
        logger.warning(
            "MarkItDown conversion failed: filename=%s mime_type=%s error=%s",
            file.filename,
            file.mime_type,
            exc,
        )
        raise MarkItDownRequestError(
            f"Failed to convert attached file '{file.filename}': {exc}",
            status_code=400,
        ) from exc

    return _normalize_markdown_text(getattr(result, "markdown", "") or "")


def _normalize_markdown_text(markdown: str) -> str:
    lines = [
        line.rstrip()
        for line in (markdown or "")
        .replace("\r\n", "\n")
        .replace("\r", "\n")
        .split("\n")
    ]
    normalized: list[str] = []
    previous_blank = False
    for line in lines:
        is_blank = not line.strip()
        if is_blank and previous_blank:
            continue
        normalized.append(line)
        previous_blank = is_blank
    return "\n".join(normalized).strip()


def _get_converter() -> Any:
    global _converter
    if _converter is not None:
        return _converter
    with _converter_lock:
        if _converter is not None:
            return _converter
        try:
            from markitdown import MarkItDown
        except ImportError as exc:
            raise RuntimeError(
                "MarkItDown is not installed. Install markitdown[pdf,docx,pptx]."
            ) from exc
        _converter = MarkItDown(enable_plugins=False)
        return _converter


def _iter_file_part_dicts(content: Any) -> list[dict[str, Any]]:
    if not isinstance(content, list):
        return []
    return [
        part
        for part in (_part_to_dict(item) for item in content)
        if part.get("type") == "file"
    ]


def _part_to_dict(part: Any) -> dict[str, Any]:
    if hasattr(part, "model_dump"):
        return part.model_dump(exclude_none=True)
    if hasattr(part, "dict"):
        return part.dict(exclude_none=True)
    return part if isinstance(part, dict) else {}


def _content_text(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return str(content)

    text_parts: list[str] = []
    for part in content:
        part_dict = _part_to_dict(part)
        if part_dict.get("type") in {"text", "input_text"}:
            text_parts.append(str(part_dict.get("text") or ""))
    return "\n".join(p for p in text_parts if p)


def _format_attachment_markdown(filename: str, markdown: str) -> str:
    return f"## Attached file: {filename}\n\n{markdown.strip()}"


def _validate_supported_file(filename: str, mime_type: str) -> None:
    extension = Path(filename).suffix.lower()
    if (
        extension in MARKITDOWN_UNSUPPORTED_SPREADSHEET_EXTENSIONS
        or mime_type in MARKITDOWN_SPREADSHEET_MIME_TYPES
    ):
        raise MarkItDownRequestError(
            "Spreadsheet attachments are not supported. "
            "Supported formats are PDF, DOCX, PPTX, TXT, and MD.",
            status_code=400,
        )
    if extension not in SUPPORTED_ATTACHMENT_EXTENSIONS:
        raise MarkItDownRequestError(
            "Unsupported attachment type. "
            "Supported formats are PDF, DOCX, PPTX, TXT, and MD.",
            status_code=400,
        )
    if mime_type and mime_type not in SUPPORTED_ATTACHMENT_MIME_TYPES:
        raise MarkItDownRequestError(
            "Unsupported attachment MIME type. "
            "Supported formats are PDF, DOCX, PPTX, TXT, and MD.",
            status_code=400,
        )


def _mime_type_for_extension(extension: str) -> str:
    if extension == ".txt":
        return "text/plain"
    if extension == ".md":
        return "text/markdown"
    if extension == ".pdf":
        return "application/pdf"
    if extension == ".docx":
        return DOCX_MIME_TYPE
    if extension == ".pptx":
        return PPTX_MIME_TYPE
    return "application/octet-stream"


def _extension_for_mime_type(mime_type: str) -> str:
    if mime_type == "application/pdf":
        return ".pdf"
    if mime_type == DOCX_MIME_TYPE:
        return ".docx"
    if mime_type == PPTX_MIME_TYPE:
        return ".pptx"
    if mime_type == "text/plain":
        return ".txt"
    if mime_type in {"text/markdown", "text/x-markdown"}:
        return ".md"
    return ""


def _is_plain_text_attachment(file: MarkItDownFile) -> bool:
    return file.extension in PLAIN_TEXT_SUPPORTED_EXTENSIONS


def _decode_data(value: str) -> bytes:
    data = value.strip()
    if data.startswith("data:"):
        marker = ";base64,"
        idx = data.find(marker)
        if idx < 0:
            raise MarkItDownRequestError("Only base64 data URIs are supported.")
        data = data[idx + len(marker) :]
    try:
        return base64.b64decode(data, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise MarkItDownRequestError("File data must be valid base64.") from exc


def _mime_type_from_data_uri(value: str) -> str:
    data = value.strip()
    if not data.startswith("data:"):
        return ""
    marker = ";base64,"
    idx = data.find(marker)
    if idx < 0:
        return ""
    return data[5:idx].strip().lower()


def _has_extractable_text(markdown: str) -> bool:
    text = markdown.strip()
    if not text:
        return False
    return any(ch.isalnum() for ch in text)


# ===== FILE: omlx/api/markitdown_pdf_fallback.py =====

# SPDX-License-Identifier: Apache-2.0
"""PDF OCR processing for the MarkItDown integration."""

from __future__ import annotations

import asyncio
import base64
import io
import logging
from typing import Any

from ..exceptions import (
    EnginePoolError,
    InsufficientMemoryError,
    ModelLoadingError,
    ModelNotFoundError,
    ModelTooLargeError,
)
from .markitdown import MarkItDownFile, MarkItDownRequestError, quiet_pdf_parser_loggers

logger = logging.getLogger(__name__)


def resolve_pdf_ocr_model(
    model_id: str,
    *,
    engine_pool: Any | None,
    settings_manager: Any | None,
) -> str:
    if engine_pool is None:
        raise MarkItDownRequestError(
            "PDF OCR processing requires an initialized engine pool.",
            status_code=503,
        )

    resolved = engine_pool.resolve_model_id(model_id, settings_manager)
    entry = engine_pool.get_entry(resolved)
    if entry is None:
        raise MarkItDownRequestError(
            f"MarkItDown PDF OCR model not found: {model_id}",
            status_code=400,
        )

    config_model_type = str(getattr(entry, "config_model_type", "") or "").lower()
    if "ocr" not in config_model_type:
        raise MarkItDownRequestError(
            "MarkItDown PDF OCR model must have OCR in config model_type: "
            f"{model_id}",
            status_code=400,
        )

    if getattr(entry, "engine_type", "") != "vlm":
        raise MarkItDownRequestError(
            f"MarkItDown PDF OCR model must be a VLM: {model_id}",
            status_code=400,
        )

    return resolved


def render_pdf_pages_to_image_data_uris(
    file: MarkItDownFile,
    *,
    resolution: int = 144,
) -> list[str]:
    quiet_pdf_parser_loggers()
    try:
        import pdfplumber
    except ImportError as exc:
        raise RuntimeError(
            "pdfplumber is not installed. Install markitdown[pdf]."
        ) from exc

    data_uris: list[str] = []
    with pdfplumber.open(io.BytesIO(file.data)) as pdf:
        for page in pdf.pages:
            try:
                image = page.to_image(
                    resolution=resolution,
                    antialias=True,
                ).original
                buffer = io.BytesIO()
                image.save(buffer, format="PNG")
                encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
                data_uris.append(f"data:image/png;base64,{encoded}")
            finally:
                page.close()
    return data_uris


async def convert_pdf_with_ocr_engine(
    file: MarkItDownFile,
    *,
    engine_model_id: str,
    engine_pool: Any | None,
    settings_manager: Any | None,
    global_settings: Any | None,
    get_sampling_params: Any | None,
) -> str:
    chunks: list[str] = []
    async for chunk in stream_pdf_with_ocr_engine(
        file,
        engine_model_id=engine_model_id,
        engine_pool=engine_pool,
        settings_manager=settings_manager,
        global_settings=global_settings,
        get_sampling_params=get_sampling_params,
    ):
        chunks.append(chunk)
    return "".join(chunks).strip()


async def stream_pdf_with_ocr_engine(
    file: MarkItDownFile,
    *,
    engine_model_id: str,
    engine_pool: Any | None,
    settings_manager: Any | None,
    global_settings: Any | None,
    get_sampling_params: Any | None,
):
    model_id = resolve_pdf_ocr_model(
        engine_model_id,
        engine_pool=engine_pool,
        settings_manager=settings_manager,
    )
    data_uris = await asyncio.to_thread(render_pdf_pages_to_image_data_uris, file)
    if not data_uris:
        raise MarkItDownRequestError(
            f"No pages found in attached PDF: {file.filename}",
            status_code=400,
        )

    logger.info(
        "Using OCR PDF processing engine: filename=%s model=%s pages=%d",
        file.filename,
        model_id,
        len(data_uris),
    )

    if get_sampling_params is None:
        raise RuntimeError("get_sampling_params callback is required for OCR PDF.")

    (
        temperature,
        top_p,
        top_k,
        repetition_penalty,
        min_p,
        presence_penalty,
        frequency_penalty,
        max_tokens,
        xtc_probability,
        xtc_threshold,
    ) = get_sampling_params(None, None, model_id)
    chat_kwargs = {
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "min_p": min_p,
        "repetition_penalty": repetition_penalty,
        "presence_penalty": presence_penalty,
        "frequency_penalty": frequency_penalty,
        "xtc_probability": xtc_probability,
        "xtc_threshold": xtc_threshold,
    }

    emitted = False
    try:
        async with engine_pool.acquire(model_id) as engine:
            async for page_number, text in _stream_pages_with_ocr(
                engine,
                data_uris,
                global_settings=global_settings,
                chat_kwargs=chat_kwargs,
            ):
                text = text.strip()
                if not text:
                    continue
                emitted = True
                yield f"### Page {page_number}\n\n{text}\n\n"
    except ModelNotFoundError as exc:
        raise MarkItDownRequestError(str(exc), status_code=404) from exc
    except ModelTooLargeError as exc:
        raise MarkItDownRequestError(str(exc), status_code=507) from exc
    except InsufficientMemoryError as exc:
        raise MarkItDownRequestError(str(exc), status_code=507) from exc
    except ModelLoadingError as exc:
        raise MarkItDownRequestError(str(exc), status_code=409) from exc
    except EnginePoolError as exc:
        raise RuntimeError(str(exc)) from exc
    finally:
        unload = getattr(engine_pool, "unload_if_idle_unpinned", None)
        if callable(unload):
            await unload(model_id)

    if not emitted:
        raise MarkItDownRequestError(
            f"OCR PDF processing produced no text for attached PDF: {file.filename}",
            status_code=400,
        )


async def _stream_pages_with_ocr(
    engine: Any,
    data_uris: list[str],
    *,
    global_settings: Any | None,
    chat_kwargs: dict[str, Any],
):
    scheduler = getattr(global_settings, "scheduler", None)
    max_concurrent = int(getattr(scheduler, "max_concurrent_requests", 1) or 1)
    semaphore = asyncio.Semaphore(max(1, max_concurrent))
    queue: asyncio.Queue[tuple[int, str, Exception | None]] = asyncio.Queue()

    async def convert_page(index: int, data_uri: str) -> None:
        try:
            async with semaphore:
                output = await engine.chat(
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "image_url",
                                    "image_url": {"url": data_uri},
                                }
                            ],
                        }
                    ],
                    **dict(chat_kwargs),
                )
                await queue.put((index, (output.text or "").strip(), None))
        except Exception as exc:
            await queue.put((index, "", exc))

    tasks = [
        asyncio.create_task(convert_page(idx, data_uri))
        for idx, data_uri in enumerate(data_uris, 1)
    ]
    pending = len(tasks)
    buffered: dict[int, str] = {}
    next_to_emit = 1

    try:
        while pending:
            page_number, text, error = await queue.get()
            pending -= 1
            if error is not None:
                for task in tasks:
                    task.cancel()
                raise error

            buffered[page_number] = text
            while next_to_emit in buffered:
                yield next_to_emit, buffered.pop(next_to_emit)
                next_to_emit += 1
    finally:
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)


# ===== FILE: omlx/api/mcp_routes.py =====

# SPDX-License-Identifier: Apache-2.0
"""
MCP (Model Context Protocol) API routes.

This module provides FastAPI routes for MCP tool management:
- GET /v1/mcp/tools - List available MCP tools
- GET /v1/mcp/servers - List MCP server status
- POST /v1/mcp/execute - Execute an MCP tool
"""

from fastapi import APIRouter, HTTPException

from .openai_models import (
    MCPExecuteRequest,
    MCPExecuteResponse,
    MCPServerInfo,
    MCPServersResponse,
    MCPToolInfo,
    MCPToolsResponse,
)

router = APIRouter(prefix="/v1/mcp", tags=["mcp"])


# Callback function to get MCP manager (set by server.py)
_get_mcp_manager = None


def set_mcp_manager_getter(getter):
    """
    Set the callback function to get MCP manager.

    Args:
        getter: A callable that returns the MCP manager instance or None
    """
    global _get_mcp_manager
    _get_mcp_manager = getter


def _get_manager():
    """Get the MCP manager instance."""
    if _get_mcp_manager is None:
        return None
    return _get_mcp_manager()


@router.get("/tools")
async def list_mcp_tools() -> MCPToolsResponse:
    """List all available MCP tools."""
    manager = _get_manager()
    if manager is None:
        return MCPToolsResponse(tools=[], count=0)

    tools = []
    for tool in manager.get_all_tools():
        tools.append(MCPToolInfo(
            name=tool.full_name,
            description=tool.description,
            server=tool.server_name,
            parameters=tool.input_schema,
        ))

    return MCPToolsResponse(tools=tools, count=len(tools))


@router.get("/servers")
async def list_mcp_servers() -> MCPServersResponse:
    """Get status of all MCP servers."""
    manager = _get_manager()
    if manager is None:
        return MCPServersResponse(servers=[])

    servers = []
    for status in manager.get_server_status():
        servers.append(MCPServerInfo(
            name=status.name,
            state=status.state.value,
            transport=status.transport.value,
            tools_count=status.tools_count,
            error=status.error,
        ))

    return MCPServersResponse(servers=servers)


@router.post("/execute")
async def execute_mcp_tool(request: MCPExecuteRequest) -> MCPExecuteResponse:
    """Execute an MCP tool."""
    manager = _get_manager()
    if manager is None:
        raise HTTPException(
            status_code=503,
            detail="MCP not configured. Start server with --mcp-config"
        )

    result = await manager.execute_tool(
        request.tool_name,
        request.arguments,
    )

    return MCPExecuteResponse(
        tool_name=result.tool_name,
        content=result.content,
        is_error=result.is_error,
        error_message=result.error_message,
    )


# ===== FILE: omlx/api/openai_models.py =====

# SPDX-License-Identifier: Apache-2.0
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
"""
Pydantic models for OpenAI-compatible API.

These models define the request and response schemas for:
- Chat completions
- Text completions
- Tool calling
- MCP (Model Context Protocol) integration
"""

import json
from typing import Any, Dict, List, Optional, Union

from pydantic import AliasChoices, BaseModel, Field, field_validator

from omlx.api.shared_models import (
    BaseUsage,
    IDPrefix,
    generate_id,
    get_unix_timestamp,
)

# =============================================================================
# Content Types
# =============================================================================


class ImageURL(BaseModel):
    """Base64 data URI for vision model input."""

    url: str  # "data:image/jpeg;base64,..."
    detail: Optional[str] = "auto"  # "low", "high", "auto"


class InputAudio(BaseModel):
    """Audio input data for multimodal models (OpenAI format)."""

    data: str  # Base64-encoded audio or data URI
    format: str = "wav"  # Audio format: wav, mp3, etc.


class FileContent(BaseModel):
    """File input for attachment preprocessing.

    ``file_data`` matches OpenAI Chat Completions file content parts.
    ``data`` is accepted as an oMLX legacy alias for dashboard clients.
    """

    filename: Optional[str] = None
    mime_type: Optional[str] = None
    file_data: Optional[str] = None
    data: Optional[str] = None
    file_id: Optional[str] = None


class ContentPart(BaseModel):
    """
    A part of a message content array.

    Supports:
    - text: Plain text content
    - image_url: Image input for vision models
    - input_audio: Audio input for multimodal audio models
    - file: Document or text input for attachment preprocessing
    """

    type: str  # "text", "image_url", "input_audio", or "file"
    text: Optional[str] = None
    image_url: Optional[ImageURL] = None
    input_audio: Optional[InputAudio] = None
    file: Optional[FileContent] = None


# =============================================================================
# Messages
# =============================================================================


class Message(BaseModel):
    """
    A message in a chat conversation.

    Supports:
    - Simple text messages (role + content string)
    - Content array messages (role + content list with text parts)
    - Tool call messages (assistant with tool_calls)
    - Tool response messages (role="tool" with tool_call_id)
    """

    role: str
    content: Optional[Union[str, List[ContentPart], List[dict]]] = None
    # Reasoning/thinking content from <think> blocks (OpenAI reasoning_content field)
    reasoning_content: Optional[str] = None
    # For assistant messages with tool calls
    tool_calls: Optional[List[dict]] = None
    # For tool response messages (role="tool")
    tool_call_id: Optional[str] = None
    # Participant name, rendered into chat template (e.g. Kimi K2/K2.5 named assistants)
    name: Optional[str] = None
    # Continue from this message instead of starting a new turn (prefill / partial mode)
    partial: bool = False

    @field_validator("tool_calls", mode="before")
    @classmethod
    def _validate_tool_call_arguments(cls, v: Any) -> Any:
        """Validate arguments on each tool_call before the raw dict is stored.

        tool_calls is typed as List[dict] for flexibility, which bypasses
        FunctionCall's own validator. Re-run the same coercion here so
        malformed arguments surface as 422 instead of crashing the chat
        template on the next turn.
        """
        if not isinstance(v, list):
            return v
        for tc in v:
            if not isinstance(tc, dict):
                continue
            func = tc.get("function")
            if not isinstance(func, dict) or "arguments" not in func:
                continue
            func["arguments"] = _coerce_tool_call_arguments(func["arguments"])
        return v


# =============================================================================
# Tool Calling
# =============================================================================


def _coerce_tool_call_arguments(v: Any) -> str:
    """Normalize a tool_call.arguments value to a JSON-object string.

    Native tool-calling chat templates (Qwen3.5/3.6, GLM-4.x, MiniMax)
    iterate `arguments.items()`, which requires the echoed value to parse
    back into a dict. Rejecting malformed inputs here turns the silent 500
    in downstream template rendering into a clear 422 that tells the client
    what to fix. Dict inputs (non-spec but common) are coerced to JSON
    strings, empty/whitespace strings normalize to ``"{}"``, and any value
    that can't round-trip into a JSON object raises ValueError.
    """
    if isinstance(v, dict):
        return json.dumps(v, ensure_ascii=False)
    if not isinstance(v, str):
        raise ValueError(
            f"arguments must be a JSON-encoded string, got {type(v).__name__}. "
            "Per the OpenAI spec tool_call.arguments is a string containing JSON, "
            'not a dict/list/number. Example: \'{"location": "Tokyo"}\'.'
        )
    stripped = v.strip()
    if not stripped:
        return "{}"
    try:
        parsed = json.loads(stripped)
    except (json.JSONDecodeError, ValueError) as e:
        snippet = stripped if len(stripped) <= 120 else stripped[:117] + "..."
        raise ValueError(
            f"arguments must be valid JSON, got parse error: {e}. "
            "This usually means the client echoed a previous tool call "
            "with a malformed arguments value. Send arguments as a "
            'JSON-encoded object string like \'{"location": "Tokyo"}\'. '
            f"Received: {snippet!r}"
        ) from e
    if not isinstance(parsed, dict):
        raise ValueError(
            f"arguments must be a JSON object, got {type(parsed).__name__}. "
            "Tool-call arguments cannot be a list, number, or bare string. "
            'Example: \'{"location": "Tokyo"}\'.'
        )
    return v


class FunctionCall(BaseModel):
    """A function call with name and arguments."""

    name: str
    arguments: str  # JSON string

    @field_validator("name", mode="before")
    @classmethod
    def _normalize_name(cls, v: Any) -> str:
        return v.strip() if isinstance(v, str) else v

    @field_validator("arguments", mode="before")
    @classmethod
    def _validate_arguments_json(cls, v: Any) -> str:
        return _coerce_tool_call_arguments(v)


class ToolCall(BaseModel):
    """A tool call from the model."""

    id: str
    type: str = "function"
    function: FunctionCall


class ToolDefinition(BaseModel):
    """Definition of a tool that can be called by the model."""

    type: str = "function"
    function: dict


# =============================================================================
# Structured Output (JSON Schema)
# =============================================================================


class ResponseFormatJsonSchema(BaseModel):
    """JSON Schema definition for structured output."""

    name: str
    description: Optional[str] = None
    schema_: dict = Field(alias="schema")  # JSON Schema specification
    strict: Optional[bool] = False

    class Config:
        populate_by_name = True


class ResponseFormat(BaseModel):
    """
    Response format specification for structured output.

    Supports:
    - "text": Default text output (no structure enforcement)
    - "json_object": Forces valid JSON output
    - "json_schema": Forces JSON matching a specific schema
    """

    type: str = "text"  # "text", "json_object", "json_schema"
    json_schema: Optional[ResponseFormatJsonSchema] = None


class StructuredOutputOptions(BaseModel):
    """vLLM-compatible structured output options.

    Exactly one field should be set. When passed via ``extra_body`` in the
    OpenAI client, the key is ``structured_outputs``.

    Supports:
    - json: JSON schema (dict or string) for logit-level enforcement
    - regex: Regular expression the output must match
    - choice: List of allowed string values (output will be exactly one)
    - grammar: EBNF/GBNF context-free grammar string
    """

    model_config = {"populate_by_name": True}

    json_schema: Optional[Union[str, dict]] = Field(None, alias="json")
    regex: Optional[str] = None
    choice: Optional[List[str]] = None
    grammar: Optional[str] = None


# =============================================================================
# Chat Completion
# =============================================================================


class StreamOptions(BaseModel):
    """Options for streaming responses."""

    include_usage: bool = False


class ChatCompletionRequest(BaseModel):
    """Request for chat completion."""

    model: str
    messages: List[Message]
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    repetition_penalty: float | None = None
    max_tokens: Optional[int] = Field(
        default=None,
        validation_alias=AliasChoices("max_tokens", "max_completion_tokens"),
    )
    stream: bool = False
    stream_options: Optional[StreamOptions] = None
    stop: Optional[List[str]] = None
    min_p: float | None = None
    xtc_probability: float | None = None
    xtc_threshold: float | None = None
    presence_penalty: float | None = None
    frequency_penalty: float | None = None
    # Tool calling
    tools: Optional[List[ToolDefinition]] = None
    tool_choice: Optional[Union[str, dict]] = None  # "auto", "none", or specific tool
    # Structured output
    response_format: Optional[Union[ResponseFormat, dict]] = None
    # vLLM-compatible structured output (grammar, regex, choice, json)
    structured_outputs: Optional[Union[StructuredOutputOptions, dict]] = None
    # vLLM/OpenAI-compatible grammar alias, normalized to structured_outputs
    guided_grammar: Optional[str] = None
    # Chat template kwargs (e.g. enable_thinking, reasoning_effort)
    chat_template_kwargs: Optional[Dict[str, Any]] = None
    # Thinking budget (max thinking tokens, None = unlimited)
    thinking_budget: Optional[int] = Field(default=None, ge=0)
    # SpecPrefill: per-request enable/disable (None = use model setting)
    specprefill: Optional[bool] = None
    # SpecPrefill: per-request keep percentage (0.1-0.5, None = use model setting)
    specprefill_keep_pct: Optional[float] = None
    # SpecPrefill: per-request threshold override (min tokens to trigger, None = use model setting)
    specprefill_threshold: Optional[int] = None
    # Seed for reproducible generation (best-effort)
    seed: Optional[int] = None

    @field_validator("stop", mode="before")
    @classmethod
    def coerce_stop(cls, v):
        """Accept stop as a single string (OpenAI compat) and wrap in a list."""
        if isinstance(v, str):
            return [v]
        return v


class AssistantMessage(BaseModel):
    """Response message from the assistant."""

    role: str = "assistant"
    content: Optional[str] = None
    reasoning_content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None


class ChatCompletionChoice(BaseModel):
    """A single choice in chat completion response."""

    index: int = 0
    message: AssistantMessage
    finish_reason: Optional[str] = "stop"


class PromptTokensDetails(BaseModel):
    """Breakdown of prompt tokens used."""

    cached_tokens: Optional[int] = None
    audio_tokens: Optional[int] = None


class Usage(BaseUsage):
    """Token usage statistics for OpenAI API.

    Extends BaseUsage with optional timing metrics (oMLX extension).
    When present, timing values are in seconds.
    """

    prompt_tokens_details: Optional[PromptTokensDetails] = None
    # Timing metrics (oMLX extension, seconds)
    model_load_duration: Optional[float] = None
    time_to_first_token: Optional[float] = None
    total_time: Optional[float] = None
    prompt_eval_duration: Optional[float] = None
    generation_duration: Optional[float] = None
    prompt_tokens_per_second: Optional[float] = None
    generation_tokens_per_second: Optional[float] = None


class ChatCompletionResponse(BaseModel):
    """Response for chat completion."""

    id: str = Field(default_factory=lambda: generate_id(IDPrefix.CHAT_COMPLETION))
    object: str = "chat.completion"
    created: int = Field(default_factory=get_unix_timestamp)
    model: str
    choices: List[ChatCompletionChoice]
    usage: Usage = Field(default_factory=Usage)


# =============================================================================
# Text Completion
# =============================================================================


class CompletionRequest(BaseModel):
    """Request for text completion."""

    model: str
    prompt: Union[str, List[str]]
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    repetition_penalty: float | None = None
    max_tokens: Optional[int] = None
    stream: bool = False
    stream_options: Optional[StreamOptions] = None
    stop: Optional[List[str]] = None
    min_p: float | None = None
    xtc_probability: float | None = None
    xtc_threshold: float | None = None
    presence_penalty: float | None = None
    frequency_penalty: float | None = None
    # Seed for reproducible generation (best-effort)
    seed: Optional[int] = None
    # Cap reasoning/thinking tokens (parity with /v1/chat/completions)
    thinking_budget: Optional[int] = Field(default=None, ge=0)

    @field_validator("stop", mode="before")
    @classmethod
    def coerce_stop(cls, v):
        """Accept stop as a single string (OpenAI compat) and wrap in a list."""
        if isinstance(v, str):
            return [v]
        return v


class CompletionChoice(BaseModel):
    """A single choice in text completion response."""

    index: int = 0
    text: str
    finish_reason: Optional[str] = "stop"


class CompletionResponse(BaseModel):
    """Response for text completion."""

    id: str = Field(default_factory=lambda: generate_id(IDPrefix.COMPLETION))
    object: str = "text_completion"
    created: int = Field(default_factory=get_unix_timestamp)
    model: str
    choices: List[CompletionChoice]
    usage: Usage = Field(default_factory=Usage)


# =============================================================================
# Models List
# =============================================================================


class ModelInfo(BaseModel):
    """Information about an available model."""

    id: str
    object: str = "model"
    created: int = Field(default_factory=get_unix_timestamp)
    owned_by: str = "omlx"
    # vLLM-compatible extension: lets OpenAI-style clients discover the
    # effective context window from the listing without a separate call
    # to /v1/models/status (see #1308).
    max_model_len: int | None = None


class ModelsResponse(BaseModel):
    """Response for listing models."""

    object: str = "list"
    data: List[ModelInfo]


# =============================================================================
# MCP (Model Context Protocol)
# =============================================================================


class MCPToolInfo(BaseModel):
    """Information about an MCP tool."""

    name: str
    description: str
    server: str
    parameters: dict = Field(default_factory=dict)


class MCPToolsResponse(BaseModel):
    """Response for listing MCP tools."""

    tools: List[MCPToolInfo]
    count: int


class MCPServerInfo(BaseModel):
    """Information about an MCP server."""

    name: str
    state: str
    transport: str
    tools_count: int
    error: Optional[str] = None


class MCPServersResponse(BaseModel):
    """Response for listing MCP servers."""

    servers: List[MCPServerInfo]


class MCPExecuteRequest(BaseModel):
    """Request to execute an MCP tool."""

    model_config = {"populate_by_name": True}

    tool_name: str = Field(validation_alias=AliasChoices("tool_name", "tool"))
    arguments: dict = Field(default_factory=dict)


class MCPExecuteResponse(BaseModel):
    """Response from executing an MCP tool."""

    tool_name: str
    content: Optional[Union[str, list, dict]] = None
    is_error: bool = False
    error_message: Optional[str] = None


# =============================================================================
# Streaming (for SSE responses)
# =============================================================================


class ChatCompletionChunkDelta(BaseModel):
    """Delta content in a streaming chunk."""

    role: Optional[str] = None
    content: Optional[str] = None
    reasoning_content: Optional[str] = None
    tool_calls: Optional[List[dict]] = None


class ChatCompletionChunkChoice(BaseModel):
    """A single choice in a streaming chunk."""

    index: int = 0
    delta: ChatCompletionChunkDelta
    finish_reason: Optional[str] = None


class ChatCompletionChunk(BaseModel):
    """A streaming chunk for chat completion."""

    id: str = Field(default_factory=lambda: generate_id(IDPrefix.CHAT_COMPLETION))
    object: str = "chat.completion.chunk"
    created: int = Field(default_factory=get_unix_timestamp)
    model: str
    choices: List[ChatCompletionChunkChoice]
    usage: Optional[Usage] = None  # Present on last chunk when include_usage=true


# ===== FILE: omlx/api/parser_tool_calls.py =====

# SPDX-License-Identifier: Apache-2.0
"""Helpers for converting parser-emitted tool calls to OpenAI models."""

import logging
import uuid

from pydantic import ValidationError

from .openai_models import FunctionCall, ToolCall

logger = logging.getLogger(__name__)


def convert_parser_tool_calls(tool_calls: list[dict] | None) -> list[ToolCall]:
    """Convert parser-emitted tool-call dicts into validated OpenAI ToolCalls.

    Parser output comes from model text and can contain malformed JSON
    arguments. Treat those as recoverable parser failures rather than letting
    Pydantic validation abort the response stream.
    """
    converted: list[ToolCall] = []
    for tool_call in tool_calls or []:
        if not isinstance(tool_call, dict):
            continue
        name = tool_call.get("name", "")
        arguments = tool_call.get("arguments", "{}") or "{}"
        try:
            converted.append(
                ToolCall(
                    id=tool_call.get("id")
                    or tool_call.get("call_id")
                    or f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=name,
                        arguments=arguments,
                    ),
                )
            )
        except (TypeError, ValueError, ValidationError) as e:
            snippet = str(arguments)
            if len(snippet) > 120:
                snippet = snippet[:117] + "..."
            logger.warning(
                "Dropping malformed parser tool call %r: %s. arguments=%r",
                name,
                e,
                snippet,
            )
            continue
    return converted


# ===== FILE: omlx/api/rerank_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Pydantic models for Cohere/Jina-compatible Rerank API.

These models define the request and response schemas for:
- /v1/rerank endpoint
"""

import uuid

from pydantic import BaseModel, Field


class RerankRequest(BaseModel):
    """
    Request for reranking documents.

    Cohere/Jina-compatible request format for the /v1/rerank endpoint.
    """

    model: str
    """ID of the model to use."""

    query: str | dict[str, str]
    """
    The search query to compare documents against. String for text-only
    rerankers. Dict with 'text' and/or 'image' (base64 data URI) for
    multimodal rerankers like Qwen3-VL-Reranker.
    """

    documents: list[str] | list[dict[str, str]]
    """
    Documents to rerank. Can be:
    - List of strings
    - List of dicts with 'text' field (and optional 'image' for multimodal
      rerankers). Image values must be base64 data URIs.
    """

    top_n: int | None = None
    """
    Number of top results to return.
    If not specified, returns all documents.
    """

    return_documents: bool = True
    """Whether to include document text in the response."""

    max_chunks_per_doc: int | None = None
    """
    Maximum chunks per document (for long documents).
    Currently not implemented.
    """


class RerankResult(BaseModel):
    """A single rerank result."""

    index: int
    """Original index of the document in the input list."""

    relevance_score: float
    """Relevance score between 0 and 1."""

    document: dict[str, str] | None = None
    """
    The document (if return_documents=True). For text-only rerankers or
    string inputs, format is {"text": "..."}. For multimodal inputs, the
    original dict (including 'image') is returned as-is.
    """


class RerankUsage(BaseModel):
    """Token usage statistics for rerank request."""

    total_tokens: int
    """Total number of tokens processed."""


class RerankResponse(BaseModel):
    """
    Response from reranking documents.

    Cohere/Jina-compatible response format for the /v1/rerank endpoint.
    """

    id: str = Field(default_factory=lambda: f"rerank-{uuid.uuid4().hex[:8]}")
    """Unique identifier for the rerank request."""

    results: list[RerankResult]
    """Reranked results sorted by relevance score (descending)."""

    model: str
    """The model used for reranking."""

    usage: RerankUsage | None = None
    """Token usage statistics."""


# ===== FILE: omlx/api/responses_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""Pydantic models for the OpenAI Responses API (/v1/responses)."""

import json
from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field, model_validator

from .shared_models import IDPrefix, generate_id, get_unix_timestamp


# =============================================================================
# Request Models
# =============================================================================


class InputItem(BaseModel):
    """A single item in the Responses API input array.

    Supports EasyInputMessage (no type field), message, function_call,
    function_call_output, and many other types from the Responses API.
    """

    # type is optional — EasyInputMessage omits it
    type: Optional[str] = None
    # message fields
    role: Optional[str] = None
    content: Optional[Union[str, List[Any]]] = None
    # function_call fields
    id: Optional[str] = None
    call_id: Optional[str] = None
    name: Optional[str] = None
    arguments: Optional[str] = None
    # function_call_output fields
    output: Optional[Union[str, List[Any], Dict[str, Any]]] = None
    # status field (present on many item types)
    status: Optional[str] = None

    model_config = {"extra": "allow"}

    @model_validator(mode="before")
    @classmethod
    def _serialize_complex_output(cls, data: Any) -> Any:
        """Serialize list/dict output to JSON string for compatibility.

        Agent frameworks may send multimodal tool outputs (e.g. images) as
        lists or dicts. Convert them to JSON strings so downstream code that
        expects ``str`` keeps working.
        """
        if isinstance(data, dict):
            output = data.get("output")
            if isinstance(output, (list, dict)):
                data = {**data, "output": json.dumps(output)}
        return data


class ResponsesTool(BaseModel):
    """Tool definition in Responses API format.

    Supports function, local_shell, mcp, web_search, and other tool types.
    """

    type: str = "function"
    # function tool fields
    name: Optional[str] = None
    description: Optional[str] = None
    parameters: Optional[Dict[str, Any]] = None
    strict: Optional[bool] = None

    model_config = {"extra": "allow"}


class TextFormatConfig(BaseModel):
    """Text format configuration."""

    type: str = "text"  # "text", "json_object", "json_schema"
    name: Optional[str] = None
    description: Optional[str] = None
    schema_: Optional[Dict[str, Any]] = Field(None, alias="schema")
    strict: Optional[bool] = None

    model_config = {"extra": "allow", "populate_by_name": True}


class TextConfig(BaseModel):
    """Text configuration wrapper."""

    format: Optional[TextFormatConfig] = None
    verbosity: Optional[str] = None  # "low", "medium", "high"

    model_config = {"extra": "allow"}


class ResponsesRequest(BaseModel):
    """Request body for POST /v1/responses."""

    model: str
    input: Optional[Union[str, List[InputItem]]] = None
    instructions: Optional[str] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    max_output_tokens: Optional[int] = None
    stream: bool = False
    tools: Optional[List[ResponsesTool]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None
    text: Optional[TextConfig] = None
    previous_response_id: Optional[str] = None
    store: Optional[bool] = None
    truncation: Optional[str] = None  # "auto" or "disabled"
    metadata: Optional[Dict[str, str]] = None
    reasoning: Optional[Dict[str, Any]] = None
    parallel_tool_calls: Optional[bool] = None
    # Fields that Codex CLI sends
    include: Optional[List[str]] = None
    service_tier: Optional[str] = None
    prompt_cache_key: Optional[str] = None
    prompt_cache_retention: Optional[str] = None
    user: Optional[str] = None
    top_logprobs: Optional[int] = None
    background: Optional[bool] = None
    conversation: Optional[Any] = None
    max_tool_calls: Optional[int] = None
    stream_options: Optional[Dict[str, Any]] = None
    # Seed for reproducible generation (best-effort)
    seed: Optional[int] = None
    # Chat template kwargs (e.g. enable_thinking, reasoning_effort)
    chat_template_kwargs: Optional[Dict[str, Any]] = None

    model_config = {"extra": "allow"}


# =============================================================================
# Response Models
# =============================================================================


class OutputContent(BaseModel):
    """Content block within an output message item."""

    type: str = "output_text"
    text: str = ""
    annotations: List[Any] = Field(default_factory=list)


class ReasoningSummaryPart(BaseModel):
    """A single part of a reasoning summary."""

    type: str = "summary_text"
    text: str = ""


class OutputItem(BaseModel):
    """A single item in the response output array.

    Can be a message, function_call, or reasoning.
    """

    type: str  # "message" or "function_call" or "reasoning"
    id: str
    status: str = "completed"
    # message fields
    role: Optional[str] = None
    content: Optional[List[OutputContent]] = None
    # function_call fields
    call_id: Optional[str] = None
    name: Optional[str] = None
    arguments: Optional[str] = None
    # reasoning fields
    summary: Optional[List[ReasoningSummaryPart]] = None


class InputTokensDetails(BaseModel):
    """Details about input token usage."""

    cached_tokens: int = 0


class OutputTokensDetails(BaseModel):
    """Details about output token usage."""

    reasoning_tokens: int = 0


class ResponseUsage(BaseModel):
    """Token usage for Responses API."""

    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    input_tokens_details: InputTokensDetails = Field(
        default_factory=InputTokensDetails
    )
    output_tokens_details: OutputTokensDetails = Field(
        default_factory=OutputTokensDetails
    )

    def model_post_init(self, __context) -> None:
        if self.total_tokens == 0 and (self.input_tokens > 0 or self.output_tokens > 0):
            object.__setattr__(
                self,
                "total_tokens",
                self.input_tokens + self.output_tokens,
            )


class ResponseObject(BaseModel):
    """Full response object for the Responses API."""

    id: str = Field(default_factory=lambda: generate_id(IDPrefix.RESPONSE))
    object: Literal["response"] = "response"
    created_at: int = Field(default_factory=get_unix_timestamp)
    model: str
    status: str = "completed"  # "completed", "in_progress", "failed", "incomplete"
    output: List[OutputItem] = Field(default_factory=list)
    usage: Optional[ResponseUsage] = None
    text: Optional[TextConfig] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
    tools: List[ResponsesTool] = Field(default_factory=list)
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    max_output_tokens: Optional[int] = None
    previous_response_id: Optional[str] = None
    metadata: Optional[Dict[str, str]] = Field(default_factory=dict)
    truncation: Optional[str] = None
    error: Optional[Dict[str, Any]] = None


# ===== FILE: omlx/api/responses_utils.py =====

# SPDX-License-Identifier: Apache-2.0
"""Conversion utilities for the OpenAI Responses API."""

import copy
import json
import logging
import uuid
from collections import OrderedDict
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from .responses_models import (
    InputItem,
    InputTokensDetails,
    OutputContent,
    OutputItem,
    ResponsesTool,
    ResponseUsage,
)
from .shared_models import IDPrefix, generate_id

logger = logging.getLogger(__name__)


class ResponseStateError(RuntimeError):
    """Base error for persisted Responses API conversation state."""


class ResponseStateNotFoundError(ResponseStateError):
    """Raised when the requested response state does not exist."""


class ResponseStateCorruptError(ResponseStateError):
    """Raised when a stored response chain is incomplete or invalid."""


def _try_parse_json(s: str):
    """Try to parse a string as JSON dict/list, return original string on failure."""
    if not isinstance(s, str):
        return s
    s = s.strip()
    if not s or not (s.startswith("{") or s.startswith("[")):
        return s
    try:
        return json.loads(s)
    except (json.JSONDecodeError, ValueError):
        return s


def _flush_pending_tool_calls(
    messages: List[Dict[str, Any]],
    pending: List[Dict[str, Any]],
    min_merge_index: int = 0,
    pending_reasoning: str = "",
) -> str:
    """Flush accumulated tool calls into messages.

    If the last message is an assistant message without tool_calls, merge
    into it (avoids duplicate assistant turns that confuse chat templates).
    Otherwise create a new assistant message.

    When ``pending_reasoning`` is set, attach it as ``reasoning_content``
    on the synthesized assistant message so reasoning round-trips even
    when the spec sequence is reasoning → function_call → output (no
    intervening message item). Returns the passthrough reasoning when
    no tool calls were flushed, or "" when reasoning was consumed.
    """
    if not pending:
        return pending_reasoning
    if (
        messages
        and len(messages) - 1 >= min_merge_index
        and messages[-1].get("role") == "assistant"
        and "tool_calls" not in messages[-1]
    ):
        messages[-1]["tool_calls"] = list(pending)
        if pending_reasoning:
            messages[-1]["reasoning_content"] = pending_reasoning
    else:
        msg: Dict[str, Any] = {"role": "assistant", "tool_calls": list(pending)}
        if pending_reasoning:
            msg["reasoning_content"] = pending_reasoning
        messages.append(msg)
    pending.clear()
    return ""


def _consolidate_system_messages(
    messages: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """Move all system messages to the front and merge them into one."""
    system_parts: List[str] = []
    non_system: List[Dict[str, Any]] = []
    for msg in messages:
        if msg.get("role") == "system":
            content = msg.get("content", "")
            if content:
                system_parts.append(content)
        else:
            non_system.append(msg)

    if not system_parts:
        return messages

    return [{"role": "system", "content": "\n\n".join(system_parts)}] + non_system


# =============================================================================
# Input Conversion
# =============================================================================


def convert_responses_input_to_messages(
    input_data: Optional[Union[str, List[InputItem]]],
    instructions: Optional[str] = None,
    previous_messages: Optional[List[Dict[str, Any]]] = None,
    consolidate_system_messages: bool = True,
) -> List[Dict[str, Any]]:
    """Convert Responses API input to internal messages format.

    Args:
        input_data: String prompt or list of InputItem objects.
        instructions: System prompt (prepended as system message).
        previous_messages: Messages from previous_response_id chain.
        consolidate_system_messages: If True, merge all system/developer content
            into one leading system message for strict templates. Server code can
            set this to False and resolve placement after the target template is
            known.

    Returns:
        List of message dicts compatible with chat template.
    """
    messages: List[Dict[str, Any]] = []

    # Collect system/developer content to merge into a single system message
    # when strict-template compatibility mode is active. In deferred mode,
    # top-level instructions still form a leading system message, but input
    # system/developer items keep their original position until template
    # capability probing decides whether they can be preserved.
    system_parts: List[str] = []
    if instructions:
        system_parts.append(instructions)

    # Prepend previous response context
    if previous_messages:
        messages.extend(copy.deepcopy(previous_messages))
    current_message_start = len(messages)

    if input_data is None:
        if system_parts:
            messages.insert(0, {"role": "system", "content": "\n\n".join(system_parts)})
        return (
            _consolidate_system_messages(messages)
            if consolidate_system_messages
            else messages
        )

    if isinstance(input_data, str):
        if system_parts:
            messages.insert(0, {"role": "system", "content": "\n\n".join(system_parts)})
        messages.append({"role": "user", "content": input_data})
        return (
            _consolidate_system_messages(messages)
            if consolidate_system_messages
            else messages
        )

    # Process input items
    # Track pending tool calls for grouping into a single assistant message
    pending_tool_calls: List[Dict[str, Any]] = []
    # Track reasoning content to attach to the next assistant message
    pending_reasoning: str = ""

    for item in input_data:
        # Resolve effective type: EasyInputMessage has no type field
        item_type = item.type
        if item_type is None and item.role is not None:
            item_type = "message"

        if item_type == "message":
            # Flush pending tool calls before a new message. Reasoning
            # passes through when no tool calls were flushed so it lands
            # on this message instead.
            pending_reasoning = _flush_pending_tool_calls(
                messages,
                pending_tool_calls,
                min_merge_index=current_message_start,
                pending_reasoning=pending_reasoning,
            )

            role = item.role or "user"
            # Map "developer" role to "system"
            if role == "developer":
                role = "system"

            content = item.content
            if isinstance(content, list):
                # Convert content parts - preserve images for VLM processing
                text_parts = []
                has_image = False
                converted_parts: List[Dict[str, Any]] = []
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") in ("input_text", "text", "output_text"):
                            text = part.get("text", "")
                            text_parts.append(text)
                            converted_parts.append({"type": "text", "text": text})
                        elif part.get("type") == "input_image":
                            # Preserve image data for VLM engines
                            has_image = True
                            image_url = part.get("image_url", part.get("url", ""))
                            detail = part.get("detail", "auto")
                            converted_parts.append(
                                {
                                    "type": "input_image",
                                    "image_url": image_url,
                                    "detail": detail,
                                }
                            )
                    elif isinstance(part, str):
                        text_parts.append(part)
                        converted_parts.append({"type": "text", "text": part})
                if has_image:
                    # Keep as content list so VLM can extract images
                    content = converted_parts
                else:
                    content = "\n".join(text_parts) if text_parts else ""

            # Merge system/developer messages into the single system block unless
            # the server is deferring placement until the template is known.
            if role == "system":
                if consolidate_system_messages:
                    system_parts.append(content or "")
                else:
                    messages.append({"role": "system", "content": content or ""})
            else:
                msg_dict: Dict[str, Any] = {"role": role, "content": content or ""}
                if role == "assistant" and pending_reasoning:
                    msg_dict["reasoning_content"] = pending_reasoning
                    pending_reasoning = ""
                messages.append(msg_dict)

        elif item_type == "reasoning":
            # Collect reasoning summary text to attach to the next
            # assistant message as reasoning_content.
            summary = getattr(item, "summary", None) or (
                (item.model_extra or {}).get("summary")
                if hasattr(item, "model_extra")
                else None
            )
            if summary:
                parts = []
                for s in summary:
                    if isinstance(s, dict):
                        parts.append(s.get("text", ""))
                    else:
                        parts.append(getattr(s, "text", ""))
                pending_reasoning = "\n".join(p for p in parts if p)

        elif item.type == "function_call":
            # Assistant's tool call — accumulate for grouping
            call_id = item.call_id or item.id or f"call_{uuid.uuid4().hex[:8]}"
            pending_tool_calls.append(
                {
                    "id": call_id,
                    "type": "function",
                    "function": {
                        "name": item.name or "",
                        "arguments": _try_parse_json(item.arguments or "{}"),
                    },
                }
            )

        elif item.type == "function_call_output":
            # Flush pending tool calls first. Any pending reasoning gets
            # attached to the synthesized assistant tool_calls message.
            pending_reasoning = _flush_pending_tool_calls(
                messages,
                pending_tool_calls,
                min_merge_index=current_message_start,
                pending_reasoning=pending_reasoning,
            )

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": item.call_id or "",
                    "content": item.output or "",
                }
            )

    # Flush remaining pending tool calls. If reasoning survived without
    # a trailing message, attach it to the synthesized tool_calls message.
    _flush_pending_tool_calls(
        messages,
        pending_tool_calls,
        min_merge_index=current_message_start,
        pending_reasoning=pending_reasoning,
    )

    # Insert merged system message at position 0
    if system_parts:
        messages.insert(0, {"role": "system", "content": "\n\n".join(system_parts)})

    return (
        _consolidate_system_messages(messages)
        if consolidate_system_messages
        else messages
    )


# =============================================================================
# Tool Conversion
# =============================================================================


def convert_responses_tools(
    tools: Optional[List[ResponsesTool]],
) -> Optional[List[Dict[str, Any]]]:
    """Convert Responses API flat tool format to Chat Completions nested format.

    Responses: {"type": "function", "name": "fn", "parameters": {...}}
    Chat Completions: {"type": "function", "function": {"name": "fn", "parameters": {...}}}

    Non-function tool types (local_shell, mcp, web_search, etc.) are skipped
    since they are not supported by local model chat templates.
    """
    if not tools:
        return None

    result = []
    for tool in tools:
        if tool.type == "function" and tool.name:
            func_def: Dict[str, Any] = {"name": tool.name}
            if tool.description:
                func_def["description"] = tool.description
            if tool.parameters:
                func_def["parameters"] = tool.parameters
            if tool.strict is not None:
                func_def["strict"] = tool.strict
            result.append({"type": "function", "function": func_def})
        # Non-function tools (local_shell, mcp, web_search, etc.) are
        # silently skipped — local models can't execute them.
    return result if result else None


# =============================================================================
# Response Building
# =============================================================================


def build_message_output_item(
    text: str,
    item_id: Optional[str] = None,
    status: str = "completed",
) -> OutputItem:
    """Build a message-type OutputItem."""
    return OutputItem(
        type="message",
        id=item_id or generate_id(IDPrefix.MESSAGE),
        status=status,
        role="assistant",
        content=[OutputContent(type="output_text", text=text)],
    )


def build_function_call_output_item(
    name: str,
    arguments: str,
    call_id: str,
    item_id: Optional[str] = None,
    status: str = "completed",
) -> OutputItem:
    """Build a function_call-type OutputItem."""
    return OutputItem(
        type="function_call",
        id=item_id or generate_id(IDPrefix.FUNCTION_CALL),
        status=status,
        call_id=call_id,
        name=name,
        arguments=arguments,
    )


def build_reasoning_output_item(
    reasoning_text: str,
    item_id: Optional[str] = None,
    status: str = "completed",
) -> OutputItem:
    """Build a reasoning-type OutputItem with full CoT in summary[0].text."""
    from .responses_models import ReasoningSummaryPart

    summary = [ReasoningSummaryPart(text=reasoning_text)] if reasoning_text else []
    return OutputItem(
        type="reasoning",
        id=item_id or generate_id(IDPrefix.REASONING),
        status=status,
        summary=summary,
    )


def build_response_usage(
    input_tokens: int,
    output_tokens: int,
    reasoning_tokens: int = 0,
    cached_tokens: int = 0,
) -> ResponseUsage:
    """Build ResponseUsage from token counts."""
    from .responses_models import OutputTokensDetails

    return ResponseUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=input_tokens + output_tokens,
        input_tokens_details=InputTokensDetails(cached_tokens=cached_tokens),
        output_tokens_details=OutputTokensDetails(reasoning_tokens=reasoning_tokens),
    )


# =============================================================================
# SSE Event Formatting
# =============================================================================


def format_sse_event(event_type: str, data: Any) -> str:
    """Format a Responses API SSE event.

    Returns: "event: {type}\\ndata: {json}\\n\\n"
    """
    if isinstance(data, str):
        json_str = data
    elif hasattr(data, "model_dump"):
        json_str = json.dumps(data.model_dump(exclude_none=True))
    elif isinstance(data, dict):
        json_str = json.dumps(data)
    else:
        json_str = json.dumps(data)
    return f"event: {event_type}\ndata: {json_str}\n\n"


# =============================================================================
# Response Store (previous_response_id support)
# =============================================================================

MAX_STORED_RESPONSES = 1000


class ResponseStore:
    """Bounded persisted store for response state and public responses."""

    def __init__(
        self,
        max_size: int = MAX_STORED_RESPONSES,
        state_dir: Optional[Union[str, Path]] = None,
    ):
        self._store: OrderedDict[str, Dict[str, Any]] = OrderedDict()
        self._max_size = max_size
        self._state_dir = Path(state_dir).expanduser().resolve() if state_dir else None
        if self._state_dir:
            self._state_dir.mkdir(parents=True, exist_ok=True)
            self._load_persisted_records()

    @property
    def state_dir(self) -> Optional[Path]:
        """Resolved directory used for persisted response state."""
        return self._state_dir

    def _record_path(self, response_id: str) -> Optional[Path]:
        if self._state_dir is None:
            return None
        return self._state_dir / f"{response_id}.json"

    def _normalize_record(
        self,
        response_id: str,
        response_data: Dict[str, Any],
    ) -> Dict[str, Any]:
        if "public_response" in response_data:
            record = copy.deepcopy(response_data)
            record.setdefault("response_id", response_id)
            record.setdefault(
                "created_at", record.get("public_response", {}).get("created_at", 0)
            )
            record.setdefault(
                "previous_response_id",
                record.get("public_response", {}).get("previous_response_id"),
            )
            record.setdefault("input_messages", [])
            record.setdefault(
                "output_messages",
                normalize_response_output_to_messages(
                    record.get("public_response", {}).get("output", [])
                ),
            )
            return record

        public_response = copy.deepcopy(response_data)
        public_response.setdefault("id", response_id)
        return {
            "response_id": response_id,
            "previous_response_id": public_response.get("previous_response_id"),
            "input_messages": [],
            "output_messages": normalize_response_output_to_messages(
                public_response.get("output", [])
            ),
            "public_response": public_response,
            "created_at": public_response.get("created_at", 0),
        }

    def _persist_record(self, record: Dict[str, Any]) -> None:
        path = self._record_path(record["response_id"])
        if path is None:
            return
        tmp_path = path.with_suffix(".tmp")
        with tmp_path.open("w", encoding="utf-8") as f:
            json.dump(record, f, ensure_ascii=False)
        tmp_path.replace(path)

    def _remove_persisted_record(self, response_id: str) -> None:
        path = self._record_path(response_id)
        if path is None or not path.exists():
            return
        path.unlink()

    def _evict_oldest(self) -> None:
        while len(self._store) > self._max_size:
            response_id, _record = self._store.popitem(last=False)
            self._remove_persisted_record(response_id)

    def _load_persisted_records(self) -> None:
        assert self._state_dir is not None
        loaded: List[Dict[str, Any]] = []
        for path in sorted(self._state_dir.glob("*.json")):
            try:
                with path.open("r", encoding="utf-8") as f:
                    raw = json.load(f)
                response_id = raw.get("response_id") or raw.get(
                    "public_response", {}
                ).get("id")
                if not response_id:
                    raise ValueError("missing response_id")
                loaded.append(self._normalize_record(response_id, raw))
            except (OSError, ValueError, json.JSONDecodeError) as exc:
                logger.warning("Skipping corrupt response state file %s: %s", path, exc)

        loaded.sort(
            key=lambda record: (record.get("created_at", 0), record["response_id"])
        )
        for record in loaded:
            self._store[record["response_id"]] = record
        self._evict_oldest()

    def put(self, response_id: str, response_data: Dict[str, Any]) -> None:
        """Store response state, evicting oldest records if needed."""
        record = self._normalize_record(response_id, response_data)
        if response_id in self._store:
            self._store.move_to_end(response_id)
        self._store[response_id] = record
        self._persist_record(record)
        self._evict_oldest()

    def get_record(self, response_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve a stored response-state record."""
        data = self._store.get(response_id)
        if data is not None:
            self._store.move_to_end(response_id)
            return copy.deepcopy(data)
        return None

    def get(self, response_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve the public response object for a stored record."""
        data = self.get_record(response_id)
        if data is None:
            return None
        return data.get("public_response")

    def resolve_chain_messages(self, response_id: str) -> List[Dict[str, Any]]:
        """Resolve the full previous_response_id chain into message history."""
        if response_id not in self._store:
            raise ResponseStateNotFoundError(f"Response state not found: {response_id}")

        chain: List[Dict[str, Any]] = []
        seen: set[str] = set()
        current_id: Optional[str] = response_id
        while current_id:
            if current_id in seen:
                raise ResponseStateCorruptError(
                    f"Cycle detected in previous_response_id chain at {current_id}"
                )
            seen.add(current_id)
            record = self._store.get(current_id)
            if record is None:
                raise ResponseStateCorruptError(
                    f"Missing ancestor response state: {current_id}"
                )
            self._store.move_to_end(current_id)
            chain.append(record)
            current_id = record.get("previous_response_id")

        chain.reverse()
        messages: List[Dict[str, Any]] = []
        for record in chain:
            messages.extend(copy.deepcopy(record.get("input_messages", [])))
            messages.extend(copy.deepcopy(record.get("output_messages", [])))
        return _consolidate_system_messages(messages)

    def delete(self, response_id: str) -> bool:
        """Delete a stored response. Returns True if found."""
        if response_id not in self._store:
            return False
        del self._store[response_id]
        self._remove_persisted_record(response_id)
        return True

    def __len__(self) -> int:
        return len(self._store)


# =============================================================================
# Previous Response Conversion
# =============================================================================


def convert_stored_response_to_messages(
    response_data: Dict[str, Any],
) -> List[Dict[str, Any]]:
    """Convert a stored public response or state record back to messages."""
    if "output_messages" in response_data:
        return copy.deepcopy(response_data.get("output_messages", []))
    return normalize_response_output_to_messages(response_data.get("output", []))


def normalize_response_output_to_messages(
    output_items: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """Convert response output items to assistant/tool-call history messages."""
    messages: List[Dict[str, Any]] = []
    pending_tool_calls: List[Dict[str, Any]] = []
    pending_reasoning: str = ""

    for item in output_items:
        item_type = item.get("type")
        if item_type == "reasoning":
            summary = item.get("summary", [])
            parts = [s.get("text", "") for s in summary if isinstance(s, dict)]
            pending_reasoning = "\n".join(p for p in parts if p)
        elif item_type == "message":
            pending_reasoning = _flush_pending_tool_calls(
                messages,
                pending_tool_calls,
                pending_reasoning=pending_reasoning,
            )
            content_blocks = item.get("content", [])
            text_parts = []
            for block in content_blocks:
                if block.get("type") == "output_text":
                    text_parts.append(block.get("text", ""))
            msg_dict: Dict[str, Any] = {
                "role": item.get("role", "assistant"),
                "content": "\n".join(text_parts),
            }
            if pending_reasoning:
                msg_dict["reasoning_content"] = pending_reasoning
                pending_reasoning = ""
            messages.append(msg_dict)
        elif item_type == "function_call":
            call_id = item.get("call_id", f"call_{uuid.uuid4().hex[:8]}")
            pending_tool_calls.append(
                {
                    "id": call_id,
                    "type": "function",
                    "function": {
                        "name": item.get("name", ""),
                        "arguments": _try_parse_json(item.get("arguments", "{}")),
                    },
                }
            )

    _flush_pending_tool_calls(
        messages,
        pending_tool_calls,
        pending_reasoning=pending_reasoning,
    )
    return _consolidate_system_messages(messages)


def build_response_store_record(
    public_response: Dict[str, Any],
    input_messages: List[Dict[str, Any]],
    output_messages: List[Dict[str, Any]],
) -> Dict[str, Any]:
    """Build a persisted response-state record."""
    return {
        "response_id": public_response.get("id", ""),
        "previous_response_id": public_response.get("previous_response_id"),
        "input_messages": copy.deepcopy(input_messages),
        "output_messages": copy.deepcopy(output_messages),
        "public_response": copy.deepcopy(public_response),
        "created_at": public_response.get("created_at", 0),
    }


# ===== FILE: omlx/api/shared_models.py =====

# SPDX-License-Identifier: Apache-2.0
"""Shared models and utilities for API responses."""

import time
import uuid
from enum import Enum

from pydantic import BaseModel


class IDPrefix(str, Enum):
    """Prefixes for generated IDs."""

    CHAT_COMPLETION = "chatcmpl"
    COMPLETION = "cmpl"
    MESSAGE = "msg"
    EMBEDDING = "emb"
    RERANK = "rerank"
    RESPONSE = "resp"
    FUNCTION_CALL = "fc"
    REASONING = "rs"


def generate_id(prefix: IDPrefix, length: int = 8) -> str:
    """Generate a unique ID with the given prefix.

    Args:
        prefix: The ID prefix to use
        length: Length of the random suffix (default 8)

    Returns:
        Generated ID string (e.g., "chatcmpl-abc12345")
    """
    if prefix == IDPrefix.MESSAGE:
        # Anthropic style: msg_<24-char-hex>
        return f"msg_{uuid.uuid4().hex[:24]}"
    if prefix == IDPrefix.RESPONSE:
        return f"resp_{uuid.uuid4().hex[:24]}"
    if prefix == IDPrefix.FUNCTION_CALL:
        return f"fc_{uuid.uuid4().hex[:8]}"
    if prefix == IDPrefix.REASONING:
        return f"rs_{uuid.uuid4().hex[:24]}"
    return f"{prefix.value}-{uuid.uuid4().hex[:length]}"


def get_unix_timestamp() -> int:
    """Get current Unix timestamp.

    Returns:
        Current time as Unix timestamp (integer seconds since epoch)
    """
    return int(time.time())


class BaseUsage(BaseModel):
    """Base class for token usage statistics.

    This provides a foundation for both OpenAI-style (prompt_tokens/completion_tokens)
    and Anthropic-style (input_tokens/output_tokens) usage tracking.
    """

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    input_tokens: int = 0
    output_tokens: int = 0

    def model_post_init(self, __context) -> None:
        """Calculate total_tokens and sync Anthropic-style aliases."""
        if self.total_tokens == 0 and (self.prompt_tokens > 0 or self.completion_tokens > 0):
            object.__setattr__(
                self,
                "total_tokens",
                self.prompt_tokens + self.completion_tokens,
            )
        if self.input_tokens == 0 and self.prompt_tokens > 0:
            object.__setattr__(self, "input_tokens", self.prompt_tokens)
        if self.output_tokens == 0 and self.completion_tokens > 0:
            object.__setattr__(self, "output_tokens", self.completion_tokens)


# ===== FILE: omlx/api/thinking.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Thinking/reasoning content parser for separating <think>...</think> blocks.

Provides both streaming (ThinkingParser) and non-streaming (extract_thinking)
interfaces for separating reasoning content from regular response content.

Used by reasoning models like DeepSeek R1, Qwen3/3.5, MiniMax that wrap
their chain-of-thought reasoning in <think>...</think> tags.
"""

import re
from collections.abc import Callable, Sequence
from typing import List, Optional, Tuple

# Tags used for thinking blocks
_OPEN_TAG = "<think>"
_CLOSE_TAG = "</think>"
_OPEN_LEN = len(_OPEN_TAG)   # 7
_CLOSE_LEN = len(_CLOSE_TAG)  # 8
_MINIMAX_OPEN_TAG = "<mm:think>"
_MINIMAX_CLOSE_TAG = "</mm:think>"
_HY3_OPEN_TAG = "<think:opensource>"
_HY3_CLOSE_TAG = "</think:opensource>"

# Regex for non-streaming extraction (complete text)
_THINKING_PATTERN = re.compile(r'<think>(.*?)</think>', re.DOTALL)
# Handle case where <think> is missing but </think> is present
# (scheduler prepends <think>\n but the tag may be split)
_THINKING_TAIL_PATTERN = re.compile(r'^(.*?)</think>', re.DOTALL)


def _safe_tokenizer_attr(tokenizer, attr: str, default=None):
    if tokenizer is None:
        return default
    try:
        return getattr(tokenizer, attr, default)
    except (AttributeError, TypeError, ValueError):
        return default


def _single_token_id(value) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _convert_token_to_id(tokenizer, token: str) -> int | None:
    convert = _safe_tokenizer_attr(tokenizer, "convert_tokens_to_ids")
    if not callable(convert):
        return None
    try:
        token_id = convert(token)
    except (AttributeError, KeyError, TypeError, ValueError):
        return None
    if token_id == _safe_tokenizer_attr(tokenizer, "unk_token_id"):
        return None
    return _single_token_id(token_id)


def _encode_prompt_ids(tokenizer, prompt: str) -> list[int] | None:
    encode = _safe_tokenizer_attr(tokenizer, "encode")
    if not callable(encode):
        return None
    try:
        return list(encode(prompt, add_special_tokens=False))
    except TypeError:
        try:
            return list(encode(prompt))
        except Exception:
            return None
    except Exception:
        return None


def _think_end_token_ids(tokenizer) -> list[int] | None:
    think_end_id = _single_token_id(_safe_tokenizer_attr(tokenizer, "think_end_id"))
    if think_end_id is not None:
        return [think_end_id]

    think_end_tag = _safe_tokenizer_attr(tokenizer, "think_end", _CLOSE_TAG)
    encoded = _encode_prompt_ids(tokenizer, think_end_tag or _CLOSE_TAG)
    if encoded:
        return encoded

    token_id = _convert_token_to_id(tokenizer, _CLOSE_TAG)
    if token_id is not None:
        return [token_id]
    return None


def prompt_opens_thinking(
    tokenizer,
    prompt: str,
    prompt_token_ids: Sequence[int] | None = None,
) -> tuple[bool, str]:
    """Return whether a raw prompt would make the engine prepend ``<think>``.

    Presentation-layer stripping must mirror the engine/scheduler decision, not
    just the raw text suffix. Some prompts can contain a literal ``<think>``
    without tokenizing to the model's think-start id, and templates can leave
    the think-start token in the final token tail without the raw string ending
    in the visible tag. When the caller already has prompt ids from the same
    tokenizer path as the scheduler, those ids are authoritative.
    """
    think_tag = (
        _safe_tokenizer_attr(tokenizer, "think_start", _OPEN_TAG) or _OPEN_TAG
    )
    if tokenizer is None:
        return prompt.rstrip().endswith(think_tag), think_tag

    think_start_id = _single_token_id(
        _safe_tokenizer_attr(tokenizer, "think_start_id")
    )
    if think_start_id is None:
        think_start_id = _convert_token_to_id(tokenizer, think_tag)
    if think_start_id is None:
        return False, think_tag

    if prompt_token_ids is None:
        prompt_ids = _encode_prompt_ids(tokenizer, prompt)
    else:
        prompt_ids = list(prompt_token_ids)
    if not prompt_ids or not think_start_id:
        return False, think_tag

    last_tokens = list(prompt_ids[-3:])
    if think_start_id not in last_tokens:
        return False, think_tag

    last_idx = len(last_tokens) - 1 - last_tokens[::-1].index(think_start_id)
    after_start = last_tokens[last_idx + 1 :]

    if after_start:
        think_end_ids = _think_end_token_ids(tokenizer)
        if think_end_ids and think_end_ids[0] in after_start:
            return False, think_tag

    return True, think_tag


def extract_thinking(text: str) -> Tuple[str, str]:
    """Extract thinking and content from complete text.

    Handles:
    - Normal: ``<think>reasoning</think>answer`` → ``("reasoning", "answer")``
    - No thinking: ``just answer`` → ``("", "just answer")``
    - Partial (no open tag): ``reasoning</think>answer`` → ``("reasoning", "answer")``
    - Empty think: ``<think></think>answer`` → ``("", "answer")``
    - Think only: ``<think>reasoning</think>`` → ``("reasoning", "")``
    - Malformed (open with no close): ``<think>everything…`` →
      ``("", "everything…")`` — recovery for V4-style models that
      occasionally skip the ``</think>`` boundary token. Without this
      fallback the entire body would be classified as thinking and the
      visible answer would be empty.

    Tag-free text is always classified as content. Mirrors
    ``ThinkingParser.finish()`` recovery semantics (`_content_emitted`
    fallback): when the model emits no thinking markers, surface the body
    as the answer so the response is never empty.

    Args:
        text: Complete model output text.

    Returns:
        Tuple of (thinking_content, regular_content).
    """
    if not text:
        return ("", "")

    text = (
        text.replace(_MINIMAX_OPEN_TAG, _OPEN_TAG)
        .replace(_MINIMAX_CLOSE_TAG, _CLOSE_TAG)
        .replace(_HY3_OPEN_TAG, _OPEN_TAG)
        .replace(_HY3_CLOSE_TAG, _CLOSE_TAG)
    )

    thinking_parts = []
    remaining = text

    # Extract all <think>...</think> blocks
    while True:
        match = _THINKING_PATTERN.search(remaining)
        if not match:
            break
        thinking_parts.append(match.group(1))
        remaining = remaining[:match.start()] + remaining[match.end():]

    if thinking_parts:
        thinking = "\n".join(thinking_parts).strip()
        return (thinking, remaining.strip())

    # Handle partial: content before </think> without <think> tag
    if '</think>' in text and '<think>' not in text:
        match = _THINKING_TAIL_PATTERN.match(text)
        if match:
            thinking = match.group(1).strip()
            remaining = text[match.end():].strip()
            return (thinking, remaining)

    # Malformed: <think> opened but never closed. Drop the open tag and
    # treat the remainder as content so the answer body is not empty.
    if '<think>' in text and '</think>' not in text:
        idx = text.index('<think>')
        before = text[:idx]
        after = text[idx + _OPEN_LEN:]
        return ("", (before + after).strip())

    return ("", text)


class ThinkingParser:
    """Stateful streaming parser for separating <think>...</think> from content.

    Handles streaming chunks where tags may span multiple chunks.
    Returns (thinking_delta, content_delta) tuples for each feed() call.

    Example::

        parser = ThinkingParser()

        # Chunk 1: "<think>Let me"
        t, c = parser.feed("<think>Let me")
        # t = "Let me", c = ""

        # Chunk 2: " think</think>Answer"
        t, c = parser.feed(" think</think>Answer")
        # t = " think", c = "Answer"

        # Flush remaining
        t, c = parser.finish()
    """

    def __init__(self, start_in_thinking: bool = False):
        self._in_thinking: bool = start_in_thinking
        self._buffer: str = ""  # Buffer for potential partial tags
        # Recovery state for malformed thinking: when the prompt prepends
        # ``<think>`` and the model never emits ``</think>`` before EOS,
        # everything we streamed went out as thinking. The streamed events
        # cannot be retracted, so finish() emits the accumulated thinking
        # text once more as content — the client will show both panels but
        # the answer body is no longer empty.
        self._close_seen: bool = False
        self._thinking_accumulated: List[str] = []
        self._content_emitted: bool = False

    def feed(self, text: str) -> Tuple[str, str]:
        """Feed a text chunk, return (thinking_delta, content_delta).

        Args:
            text: New text chunk from model output.

        Returns:
            Tuple of (thinking_text, content_text) extracted from this chunk.
        """
        if not text:
            return ("", "")

        # Prepend any buffered partial tag content
        text = self._buffer + text
        self._buffer = ""

        thinking_out = []
        content_out = []

        i = 0
        while i < len(text):
            if text[i] == '<':
                # Check if this could be a tag start
                remaining = text[i:]

                # Try to match <think>
                if remaining.startswith(_OPEN_TAG):
                    self._in_thinking = True
                    i += _OPEN_LEN
                    continue

                if remaining.startswith(_HY3_OPEN_TAG):
                    self._in_thinking = True
                    i += len(_HY3_OPEN_TAG)
                    continue

                # Try to match </think>
                if remaining.startswith(_CLOSE_TAG):
                    self._in_thinking = False
                    self._close_seen = True
                    i += _CLOSE_LEN
                    continue

                if remaining.startswith(_HY3_CLOSE_TAG):
                    self._in_thinking = False
                    self._close_seen = True
                    i += len(_HY3_CLOSE_TAG)
                    continue

                # Check if it could be a partial tag (not enough chars yet)
                if self._could_be_tag(remaining):
                    # Buffer the rest and wait for more data
                    self._buffer = remaining
                    break

                # Not a tag, emit the '<' as regular content
                if self._in_thinking:
                    thinking_out.append('<')
                else:
                    content_out.append('<')
                i += 1
            else:
                if self._in_thinking:
                    thinking_out.append(text[i])
                else:
                    content_out.append(text[i])
                i += 1

        thinking_delta = "".join(thinking_out)
        content_delta = "".join(content_out)
        if thinking_delta:
            self._thinking_accumulated.append(thinking_delta)
        if content_delta:
            self._content_emitted = True
        return (thinking_delta, content_delta)

    def finish(self) -> Tuple[str, str]:
        """Flush any remaining buffered content.

        Should be called when the stream is complete to emit any
        buffered characters that were waiting for potential tag completion.
        Also recovers from malformed thinking — when the model never
        emitted ``</think>`` and no content was ever produced, returns
        the accumulated thinking text as content so the client surfaces
        a non-empty answer body.

        Returns:
            Tuple of (thinking_text, content_text) from remaining buffer
            (plus recovered content if applicable).
        """
        partial = self._buffer
        self._buffer = ""

        # Recovery: prompt opened a thinking block (or model echoed
        # ``<think>`` itself), the close tag never arrived, and nothing
        # ever streamed as content. Re-emit the accumulated thinking text
        # as content so the answer body is not empty. The thinking events
        # already streamed live cannot be retracted, so the client sees
        # the same text twice — once in the thinking panel, once as the
        # answer. UX trade-off documented in the chat template plan.
        if (
            self._in_thinking
            and not self._close_seen
            and not self._content_emitted
            and self._thinking_accumulated
        ):
            recovered = "".join(self._thinking_accumulated) + partial
            self._content_emitted = True
            return ("", recovered)

        if not partial:
            return ("", "")

        # Partial tag never completed — emit it as-is in the current mode.
        if self._in_thinking:
            self._thinking_accumulated.append(partial)
            return (partial, "")
        else:
            self._content_emitted = True
            return ("", partial)

    @staticmethod
    def _could_be_tag(text: str) -> bool:
        """Check if text could be the start of a <think> or </think> tag.

        Returns True if text is a proper prefix of either tag but not
        yet a complete match.
        """
        length = len(text)
        if length >= len(_HY3_CLOSE_TAG):
            # Long enough to determine - not a partial tag
            return False

        # Check against all recognised tags
        for tag in (_OPEN_TAG, _CLOSE_TAG, _HY3_OPEN_TAG, _HY3_CLOSE_TAG):
            if length < len(tag) and tag[:length] == text:
                return True

        return False


class ThinkingBudgetProcessor:
    """Logits processor that enforces a thinking token budget.

    Counts tokens generated while in thinking mode.  When the budget is
    exceeded, forces the close-think token(s) one at a time, then becomes
    a no-op for the rest of generation.

    Handles both single-token and multi-token close-think sequences, and
    supports alternative think markers (e.g. ``<longcat_think>``).

    Args:
        think_end_token_ids: Token ID(s) for the close-think tag.
        budget: Maximum number of thinking tokens before forcing close.
        think_start_token_id: Token ID for the open-think tag (re-entry detection).
    """

    def __init__(
        self,
        think_end_token_ids: List[int],
        budget: int,
        think_start_token_id: Optional[int] = None,
        leading_token_ids: Optional[List[int]] = None,
        trailing_token_ids: Optional[List[int]] = None,
        token_to_piece: Optional[Callable[[int], str | bytes | None]] = None,
    ):
        self._think_end_ids = think_end_token_ids
        # Full force sequence: \n + </think> + \n\n (matches training pattern)
        self._force_sequence = (
            (leading_token_ids or [])
            + list(think_end_token_ids)
            + (trailing_token_ids or [])
        )
        self._budget = budget
        self._think_start_id = think_start_token_id
        self._token_to_piece = token_to_piece

        # State
        self._thinking_tokens: int = 0
        self._in_thinking: bool = True  # Starts True (prompt ends with <think>)
        self._forcing: bool = False
        self._waiting_utf8: bool = False
        self._force_idx: int = 0
        self._done: bool = False
        self._first_call: bool = True
        # Sliding window for multi-token end detection
        self._recent_tokens: List[int] = []
        self._last_token_utf8_complete: bool = True
        self._pending_utf8: bytes = b""

    def __call__(self, tokens, logits):
        """mlx-lm logits processor: (tokens, logits) -> logits."""
        import mlx.core as mx

        # In new mlx-lm API, tokens is the full history list.
        # Accept each genuinely generated token exactly once (see grammar.py).
        n = len(tokens)
        if not hasattr(self, "_accepted_up_to"):
            self._accepted_up_to = n  # skip prompt tokens
        elif n > self._accepted_up_to:
            for i in range(self._accepted_up_to, n):
                self._update_state(int(tokens[i]))
            self._accepted_up_to = n

        # If state changed by _update_state, handle immediately
        if self._done:
            return logits

        if self._forcing:
            return self._force_next_token(logits, mx)

        if self._waiting_utf8:
            return logits

        if self._in_thinking:
            self._thinking_tokens += 1
            if self._thinking_tokens >= self._budget:
                if self._last_token_utf8_complete:
                    self._forcing = True
                    self._force_idx = 0
                    self._recent_tokens = []
                    return self._force_next_token(logits, mx)
                self._waiting_utf8 = True
                self._recent_tokens = []

        return logits

    def _update_state(self, token_id: int) -> None:
        """Update thinking state based on the last generated token."""
        self._last_token_utf8_complete = self._is_utf8_complete(token_id)

        if self._done:
            if self._think_start_id and token_id == self._think_start_id:
                self._in_thinking = True
                self._done = False
                self._thinking_tokens = 0
                self._recent_tokens = []
            return

        if self._forcing:
            self._force_idx += 1
            if self._force_idx >= len(self._force_sequence):
                self._in_thinking = False
                self._forcing = False
                self._done = True
            return

        # Detect natural close-think via sliding window
        if len(self._think_end_ids) == 1:
            if token_id == self._think_end_ids[0]:
                self._in_thinking = False
                self._done = True
                return
        else:
            self._recent_tokens.append(token_id)
            if len(self._recent_tokens) > len(self._think_end_ids):
                self._recent_tokens.pop(0)
            if self._recent_tokens == self._think_end_ids:
                self._in_thinking = False
                self._done = True
                return

        if self._waiting_utf8:
            if self._last_token_utf8_complete:
                self._waiting_utf8 = False
                self._forcing = True
                self._force_idx = 0
                self._recent_tokens = []
            return

        # Detect re-entry into thinking (rare but possible)
        if not self._in_thinking and self._think_start_id and token_id == self._think_start_id:
            self._in_thinking = True
            self._done = False
            self._thinking_tokens = 0
            self._recent_tokens = []

    def _is_utf8_complete(self, token_id: int) -> bool:
        """Best-effort UTF-8 boundary check for accepted token bytes."""
        if self._token_to_piece is None:
            return True
        try:
            piece = self._token_to_piece(token_id)
        except Exception:
            return True
        if piece is None:
            return True
        if isinstance(piece, str):
            self._pending_utf8 = b""
            return True
        self._pending_utf8 += piece
        try:
            self._pending_utf8.decode("utf-8")
            self._pending_utf8 = b""
            return True
        except UnicodeDecodeError as exc:
            if exc.reason == "unexpected end of data" or exc.end == len(
                self._pending_utf8
            ):
                return False
            self._pending_utf8 = b""
            return True

    def _force_next_token(self, logits, mx):
        """Force the next token in the close-think + trailing sequence."""
        target_id = self._force_sequence[self._force_idx]
        forced = mx.full(logits.shape, float("-inf"))
        forced[..., target_id] = 0.0
        return forced


# ===== FILE: omlx/api/tool_calling.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Tool calling parsing and conversion utilities.

Uses mlx-lm's modular tool parser system to support multiple model formats:
- json_tools: Pure JSON format
- minimax_m2: MiniMax M2 XML format
- function_gemma: Google Gemma function calling format
- glm47: GLM-4.7 format
- qwen3_coder: Qwen3 Coder XML format

The tool parser is automatically selected based on the model's chat template.

Also includes structured output (JSON Schema) utilities:
- parse_json_output: Extract JSON from model output
- validate_json_schema: Validate JSON against a schema
"""

import ast
import bisect
import json
import logging
import re
import uuid
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import regex
from jsonschema import ValidationError, validate

from .openai_models import FunctionCall, ResponseFormat, ToolCall

logger = logging.getLogger(__name__)


def _template_safe_description(value: Any) -> str:
    """Return a string description safe for strict chat templates."""
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    return str(value)


def _copy_schema_with_template_defaults(value: Any, *, is_schema: bool) -> Any:
    """Copy JSON Schema data while filling missing schema descriptions."""
    if isinstance(value, dict):
        copied = {}
        for key, child in value.items():
            if key == "properties" and isinstance(child, dict):
                copied[key] = {
                    name: _copy_schema_with_template_defaults(
                        prop_schema, is_schema=True
                    )
                    for name, prop_schema in child.items()
                }
            elif key in {
                "items",
                "additionalProperties",
                "contains",
                "propertyNames",
                "not",
                "if",
                "then",
                "else",
            }:
                copied[key] = _copy_schema_with_template_defaults(child, is_schema=True)
            elif key in {"oneOf", "anyOf", "allOf", "prefixItems"} and isinstance(
                child, list
            ):
                copied[key] = [
                    _copy_schema_with_template_defaults(item, is_schema=True)
                    for item in child
                ]
            else:
                copied[key] = _copy_schema_with_template_defaults(
                    child, is_schema=False
                )

        if is_schema:
            copied["description"] = _template_safe_description(
                copied.get("description")
            )
        return copied

    if isinstance(value, list):
        return [
            _copy_schema_with_template_defaults(item, is_schema=False) for item in value
        ]

    return value


def _serialize_tool_call_arguments(arguments: Any) -> str:
    """Serialize parser output to a JSON-object arguments string.

    Chat templates for models with native tool calling (Qwen 3.5/3.6 XML,
    GLM, MiniMax) iterate `arguments.items()` when the call is echoed back
    in history. Anything that does not represent a JSON object must be
    coerced to "{}" here so we never hand the client a non-JSON value that
    the next turn's template would crash on.
    """
    if isinstance(arguments, dict):
        return json.dumps(arguments, ensure_ascii=False)
    # mlx-vlm / mlx-lm gemma4 parser returns a JSON-object string per the
    # OpenAI spec. Accept it when it parses back to a dict.
    if isinstance(arguments, str):
        try:
            parsed = json.loads(arguments)
        except (json.JSONDecodeError, ValueError):
            parsed = None
        if isinstance(parsed, dict):
            return json.dumps(parsed, ensure_ascii=False)
    logger.warning(
        "Tool parser returned non-dict arguments (type=%s, repr=%.200r); "
        "coercing to empty object to keep downstream template safe.",
        type(arguments).__name__,
        arguments,
    )
    return "{}"


@dataclass(frozen=True)
class ToolCallExtraction:
    """Parsed tool-call result plus sanitized reasoning text."""

    cleaned_text: str
    tool_calls: Optional[List[ToolCall]]
    cleaned_thinking: str
    tool_calls_from_thinking: bool = False


# Declared-type buckets for schema-aware parameter coercion, mirroring
# mlx-lm's qwen3_coder parser so fallback-recovered calls keep the same
# argument types as natively parsed calls (#2332).
_SCHEMA_STRING_TYPES = {"string", "str", "text", "varchar", "char", "enum"}
_SCHEMA_BOOL_TYPES = {"boolean", "bool", "binary"}
_SCHEMA_CONTAINER_TYPES = {"object", "array", "arr"}
_SCHEMA_INT_PREFIXES = ("int", "uint", "long", "short", "unsigned")


def _tool_param_properties(func_name: str, tools: Optional[List]) -> dict:
    """Return the declared parameter properties for a tool, or {}."""
    if not tools:
        return {}
    for tool in tools:
        if not isinstance(tool, dict):
            continue
        func = tool.get("function")
        if not isinstance(func, dict) or func.get("name") != func_name:
            continue
        params = func.get("parameters")
        if isinstance(params, dict):
            props = params.get("properties")
            if isinstance(props, dict):
                return props
        return {}
    return {}


def _repair_json_value(val: str) -> Optional[Any]:
    """Best-effort repair of near-valid JSON with unbalanced brackets.

    Rewrites closing brackets that do not match the innermost open bracket
    (a common small-model slip, e.g. closing an array with "}"), drops
    closers with no matching opener, terminates an unclosed string, and
    appends missing closers.  Returns the parsed value, or None when the
    repaired text still fails to parse.
    """
    out: List[str] = []
    stack: List[str] = []
    in_string = False
    escaped = False
    for ch in val:
        if in_string:
            out.append(ch)
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False
            continue
        if ch == '"':
            in_string = True
            out.append(ch)
        elif ch in "[{":
            stack.append("]" if ch == "[" else "}")
            out.append(ch)
        elif ch in "]}":
            if stack:
                out.append(stack.pop())
        else:
            out.append(ch)
    if in_string:
        out.append('"')
    while stack:
        out.append(stack.pop())
    try:
        return json.loads("".join(out), strict=False)
    except (json.JSONDecodeError, ValueError):
        return None


def _coerce_param_value(val: str, key: str, props: dict, func_name: str) -> Any:
    """Convert an XML-extracted parameter value per its declared schema type.

    Mirrors the type conversion mlx-lm's native qwen3_coder parser applies,
    with an extra bracket-repair pass for declared container params whose
    value is near-valid JSON (#2332).  Parameters without a usable declared
    type keep the legacy best-effort JSON parse.
    """
    spec = props.get(key)
    raw_type = spec.get("type") if isinstance(spec, dict) else None
    if not isinstance(raw_type, str):
        # Undeclared param, union type list, or anyOf: legacy behavior.
        try:
            return json.loads(val)
        except (json.JSONDecodeError, ValueError):
            return val
    if val.strip().lower() == "null":
        return None
    ptype = raw_type.strip().lower()
    if ptype in _SCHEMA_STRING_TYPES:
        # A JSON-quoted string literal (e.g. MiniMax emits "SF" for a string
        # param) carries JSON encoding on the wire; decode it back to the
        # underlying string.  Plain values that merely look like JSON (42,
        # {"a": 1}) are kept verbatim so a string param never gets coerced to a
        # non-string.
        stripped = val.strip()
        if len(stripped) >= 2 and stripped[0] == '"' and stripped[-1] == '"':
            try:
                decoded = json.loads(stripped)
            except (json.JSONDecodeError, ValueError):
                decoded = None
            if isinstance(decoded, str):
                return decoded
        return val
    if ptype in _SCHEMA_BOOL_TYPES:
        lowered = val.strip().lower()
        if lowered in ("true", "false"):
            return lowered == "true"
    if ptype.startswith(_SCHEMA_INT_PREFIXES):
        try:
            return int(val.strip())
        except ValueError:
            pass
    elif ptype.startswith(("num", "float")):
        try:
            num = float(val.strip())
            return int(num) if num == int(num) else num
        except (ValueError, OverflowError):
            pass
    try:
        return json.loads(val, strict=False)
    except (json.JSONDecodeError, ValueError):
        pass
    try:
        literal = ast.literal_eval(val)
        if isinstance(literal, (dict, list, tuple)):
            return list(literal) if isinstance(literal, tuple) else literal
    except (ValueError, SyntaxError, TypeError, MemoryError):
        pass
    if ptype in _SCHEMA_CONTAINER_TYPES or ptype.startswith(("dict", "list")):
        repaired = _repair_json_value(val)
        if repaired is not None:
            logger.warning(
                "Repaired malformed JSON for parameter %r of tool %r "
                "(declared type %r)",
                key,
                func_name,
                ptype,
            )
            return repaired
        logger.warning(
            "Parameter %r of tool %r failed to parse as declared type %r; "
            "keeping raw string",
            key,
            func_name,
            ptype,
        )
    return val


def _parse_xml_tool_calls(
    text: str, tools: Optional[List] = None
) -> Tuple[str, Optional[List[ToolCall]]]:
    """
    Fallback parser for XML-based tool call formats.

    Handles models that use <tool_call>...</tool_call> XML format, including:
    - GLM format: <tool_call>func<arg_key>k</arg_key><arg_value>v</arg_value></tool_call>
    - Qwen/Llama format: <tool_call><function=name><parameter=key>value</parameter></function></tool_call>
    - Generic JSON: <tool_call>{"name": ..., "arguments": ...}</tool_call>

    When ``tools`` is provided, parameter values are coerced to their
    declared schema types instead of best-effort JSON parsing.

    Returns:
        Tuple of (cleaned_text, tool_calls or None)
    """
    tool_calls = []
    pattern = r"<tool_call>(.*?)</tool_call>"
    matches = re.findall(pattern, text, re.DOTALL)

    for match in matches:
        content = match.strip()
        try:
            # Try JSON format first: {"name": "func", "arguments": {...}}
            parsed = json.loads(content, strict=False)
            name = parsed.get("name", "")
            arguments = parsed.get("arguments", {})
            tool_calls.append(
                ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=name,
                        arguments=_serialize_tool_call_arguments(arguments),
                    ),
                )
            )
            continue
        except (json.JSONDecodeError, AttributeError):
            pass

        # Qwen/Llama format: <function=name><parameter=key>value</parameter></function>
        func_match = re.match(r"<function=(\w+)>(.*?)</function>", content, re.DOTALL)
        if func_match:
            func_name = func_match.group(1)
            params_text = func_match.group(2)
            props = _tool_param_properties(func_name, tools)
            arguments = {}
            for pm in re.finditer(
                r"<parameter=(\w+)>\s*(.*?)\s*</parameter>", params_text, re.DOTALL
            ):
                key = pm.group(1)
                val = pm.group(2).strip()
                arguments[key] = _coerce_param_value(val, key, props, func_name)
            tool_calls.append(
                ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=func_name,
                        arguments=json.dumps(arguments, ensure_ascii=False),
                    ),
                )
            )
            continue

        # GLM XML format: func_name<arg_key>k</arg_key><arg_value>v</arg_value>...
        arg_keys = re.findall(r"<arg_key>(.*?)</arg_key>", content)
        arg_values = re.findall(r"<arg_value>(.*?)</arg_value>", content, re.DOTALL)
        if arg_keys:
            # Function name is the text before the first <arg_key>
            name_match = re.match(r"^(.*?)<arg_key>", content, re.DOTALL)
            func_name = (
                name_match.group(1).strip()
                if name_match
                else content.split("<")[0].strip()
            )
            props = _tool_param_properties(func_name, tools)
            arguments = {}
            for k, v in zip(arg_keys, arg_values):
                arguments[k] = _coerce_param_value(v, k, props, func_name)
            tool_calls.append(
                ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=func_name,
                        arguments=json.dumps(arguments, ensure_ascii=False),
                    ),
                )
            )

    if not tool_calls:
        return text, None

    # Remove tool call tags from text
    cleaned = re.sub(r"<tool_call>.*?</tool_call>", "", text, flags=re.DOTALL).strip()
    return cleaned, tool_calls


def _parse_namespaced_tool_calls(
    text: str, namespace: str, tools: Optional[List] = None
) -> Tuple[str, Optional[List[ToolCall]]]:
    """
    Parse namespaced tool call tags like <minimax:tool_call>...</minimax:tool_call>.

    Handles the <invoke name="func"><parameter name="key">value</parameter></invoke>
    format used by MiniMax and similar models.

    When ``tools`` is provided, parameter values are coerced to their
    declared schema types instead of best-effort JSON parsing.

    Returns:
        Tuple of (cleaned_text, tool_calls or None)
    """
    tool_calls = []
    tag_start = f"<{namespace}:tool_call>"
    tag_end = f"</{namespace}:tool_call>"
    pattern = re.escape(tag_start) + r"(.*?)" + re.escape(tag_end)
    matches = re.findall(pattern, text, re.DOTALL)

    for match in matches:
        content = match.strip()
        # Parse <invoke name="func_name">...<parameter name="key">value</parameter>...</invoke>
        for invoke_match in re.finditer(
            r'<invoke\s+name="([^"]+)">(.*?)</invoke>', content, re.DOTALL
        ):
            func_name = invoke_match.group(1)
            params_text = invoke_match.group(2)
            props = _tool_param_properties(func_name, tools)
            arguments = {}
            for pm in re.finditer(
                r'<parameter\s+name="([^"]+)">(.*?)</parameter>', params_text, re.DOTALL
            ):
                key = pm.group(1)
                val = pm.group(2).strip()
                arguments[key] = _coerce_param_value(val, key, props, func_name)
            tool_calls.append(
                ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=func_name,
                        arguments=json.dumps(arguments, ensure_ascii=False),
                    ),
                )
            )

    if not tool_calls:
        return text, None

    cleaned = re.sub(pattern, "", text, flags=re.DOTALL).strip()
    return cleaned, tool_calls


def _parse_hermes_tool_calls(text: str) -> Tuple[str, Optional[List[ToolCall]]]:
    """
    Fallback parser for Hermes-style tool call formats.

    Handles outputs that use <|tool_call_start|>...<|tool_call_end|> markers
    with bracket-style content inside:
        <|tool_call_start|>[function_name(arg1=value1, arg2=value2)]<|tool_call_end|>

    Also handles JSON variant:
        <|tool_call_start|>{"name": "func", "arguments": {...}}<|tool_call_end|>

    Some clients/agents emit tool calls using this Hermes-style wire format.

    Returns:
        Tuple of (cleaned_text, tool_calls or None)
    """
    tool_calls = []
    pattern = r"<\|tool_call_start\|>(.*?)<\|tool_call_end\|>"
    matches = re.findall(pattern, text, re.DOTALL)

    for match in matches:
        content = match.strip()

        # Try JSON format first: {"name": "func", "arguments": {...}}
        try:
            parsed = json.loads(content)
            name = parsed.get("name", "")
            arguments = parsed.get("arguments", {})
            if name:
                tool_calls.append(
                    ToolCall(
                        id=f"call_{uuid.uuid4().hex[:8]}",
                        type="function",
                        function=FunctionCall(
                            name=name,
                            arguments=_serialize_tool_call_arguments(arguments),
                        ),
                    )
                )
                continue
        except (json.JSONDecodeError, AttributeError):
            pass

        # Hermes bracket format: [func_name(arg1=val1), other_tool(arg2=val2)]
        # The payload is Python-expression-like; use ast so commas inside quoted
        # strings or nested lists/dicts do not split calls incorrectly.
        try:
            parsed_expr = ast.parse(content, mode="eval").body
        except SyntaxError:
            parsed_expr = None

        calls = parsed_expr.elts if isinstance(parsed_expr, ast.List) else [parsed_expr]
        for call in calls:
            if not isinstance(call, ast.Call):
                continue

            if isinstance(call.func, ast.Name):
                func_name = call.func.id
            elif isinstance(call.func, ast.Attribute):
                func_name = ast.unparse(call.func)
            else:
                continue

            arguments = {}
            for kw in call.keywords:
                if kw.arg is None:
                    continue
                try:
                    arguments[kw.arg] = ast.literal_eval(kw.value)
                except (ValueError, SyntaxError):
                    arguments[kw.arg] = ast.unparse(kw.value)

            tool_calls.append(
                ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    type="function",
                    function=FunctionCall(
                        name=func_name,
                        arguments=json.dumps(arguments, ensure_ascii=False),
                    ),
                )
            )

    if not tool_calls:
        return text, None

    cleaned = re.sub(pattern, "", text, flags=re.DOTALL).strip()
    return cleaned, tool_calls


def _parse_bracket_tool_calls(text: str) -> Tuple[str, Optional[List[ToolCall]]]:
    """
    Fallback parser for bracket-style tool call formats.

    Recognizes both ``[Calling tool: name(args)]`` and ``[Tool call: name(args)]``
    prefixes, with or without arguments.  Models may emit the args-less form
    ``[Tool call: name]`` when mimicking conversation history.

    Returns:
        Tuple of (cleaned_text, tool_calls or None)
    """
    tool_calls = []
    # Match with args first (higher fidelity)
    pattern_with_args = (
        r"\[(?:Calling tool|Tool call):\s*([A-Za-z_][\w.-]*)\(({.*?})\)\]"
    )
    matched_spans: list = []
    for match in re.finditer(pattern_with_args, text, re.DOTALL):
        name = match.group(1)
        args_str = match.group(2)
        try:
            arguments = json.loads(args_str)
        except (json.JSONDecodeError, ValueError):
            arguments = {"raw": args_str}
        tool_calls.append(
            ToolCall(
                id=f"call_{uuid.uuid4().hex[:8]}",
                type="function",
                function=FunctionCall(
                    name=name,
                    arguments=json.dumps(arguments, ensure_ascii=False),
                ),
            )
        )
        matched_spans.append(match.span())

    # Match without args (model-generated simplified form)
    pattern_no_args = r"\[(?:Calling tool|Tool call):\s*([A-Za-z_][\w.-]*)\]"
    for match in re.finditer(pattern_no_args, text):
        # Skip if this span overlaps with an already-matched with-args span
        start, end = match.span()
        if any(s <= start < e for s, e in matched_spans):
            continue
        name = match.group(1)
        tool_calls.append(
            ToolCall(
                id=f"call_{uuid.uuid4().hex[:8]}",
                type="function",
                function=FunctionCall(
                    name=name,
                    arguments="{}",
                ),
            )
        )
        matched_spans.append((start, end))

    if not tool_calls:
        return text, None

    # Remove all matched spans from text
    cleaned = re.sub(pattern_with_args, "", text, flags=re.DOTALL)
    cleaned = re.sub(pattern_no_args, "", cleaned).strip()
    return cleaned, tool_calls


# ---------------------------------------------------------------------------
# Gemma 4 robust fallback parser
# ---------------------------------------------------------------------------

# Gemma 4's non-standard string delimiter (mlx_lm.tool_parsers.gemma4 uses
# the same literal in its regex).
_GEMMA4_STR_DELIM = '<|"|>'

# Bounds for parsing model-emitted arguments.  Model output is untrusted and
# attacker-influenceable (prompt injection can steer emissions verbatim), so
# parsing must stay linear-time and bounded: breaching a bound is a clean
# parse failure that flows into the existing drop-with-warning path, never an
# exception that escapes the parse chain.
_GEMMA4_MAX_ARGS_LEN = 262_144
_GEMMA4_MAX_DEPTH = 64


class _Gemma4ArgsTooComplexError(ValueError):
    """A defensive bound (length/depth) was breached parsing args.

    Distinct from an ordinary parse failure so the orchestrator can reject
    hard rather than retry with the legacy parser: the bounds are DoS guards
    against attacker-influenceable model output, and the legacy parser would
    happily parse oversized/deeply-nested input and defeat them.  Subclasses
    ValueError so the public parse chain still treats it as a clean drop.
    """

# A tool-call head: the name plus its opening ``{``.  Only the head is matched
# by regex; the argument span is found by _scan_gemma4_args_span, not by a
# recursive pattern (see that function).  The name segment captures namespaced
# MCP names (colon/dot/hyphen separated, e.g.
# call:google:mcp:text_generation:create-pdf-file, #1830).  The ``call:``
# opener is made optional and tolerant — ``(?:call)?:?`` — so the diffusion
# lane's degenerate prefixes (``calldone{`` missing the colon, ``:done{``
# missing ``call``, #1837) still match; the fallback only runs on
# marker-delimited content, so a permissive prefix cannot misfire on prose.
#
# Compiled with the ``regex`` module, NOT ``re``: once the ``call:`` literal
# anchor became optional (above), ``re``'s engine restarts the greedy
# ``[\w.-]+`` match at every position of a long bare argument value and
# backtracks O(n^2) hunting an opening ``{`` that never comes, hanging on
# adversarial output (a 300 KB bare value pegs a core indefinitely).  The
# ``regex`` engine fails that same partial match fast, so finditer stays
# linear.  See test_oversized_args_fail_cleanly.
_GEMMA4_CALL_HEAD = regex.compile(r"(?:call)?:?([\w.-]+(?::[\w.-]+)*)\{")

# The parenthesized variant of the head (#1846): under instruction-dense
# agentic load (large system prompt + many tool schemas + a big tool result)
# Gemma 4 26B reproducibly degrades from ``call:name{...}`` to a Python-kwargs
# shell ``call:name(key=value, ...)`` — same name grammar and same ``<|"|>``
# string quoting, only the OUTER ``{}`` becomes ``()`` and the top-level
# ``:`` separator becomes ``=``.  The nested grammar (objects, arrays, strings)
# is unchanged, so the hardened transcoder is reused for the inner content; only
# the outer shell needs a separate head + span scan.  Compiled with ``regex``,
# not ``re``, for the SAME reason as ``_GEMMA4_CALL_HEAD``: the optional/tolerant
# prefix makes ``re`` backtrack O(n^2) on a long bare value hunting an opening
# ``(`` that never comes (the #1854 ReDoS).  See test_oversized_paren_args.
_GEMMA4_CALL_HEAD_PAREN = regex.compile(r"(?:call)?:?([\w.-]+(?::[\w.-]+)*)\(")

# Matching close character for each balanced opener the parsers track.
_GEMMA4_CLOSE_CHAR = {"{": "}", "[": "]", "(": ")"}


def _squote_close_positions(s: str) -> list:
    """Indices of single quotes that can CLOSE a single-quoted value.

    A closing quote is one whose next non-whitespace character is ``,``,
    ``}``, ``]``, ``)`` or end of input.  ``)`` is an anchor for the
    parenthesized variant (#1846): a single-quoted value can be the last
    argument right before the call's closing paren, as in
    ``call:f(msg='hi')`` where the quote is followed immediately by ``)``.
    The curly form never closes a value with ``)`` (its values end at
    ``,``/``}``/``]``), so adding it does not change curly parsing.
    Anchoring closes this way (rather than taking the first quote) keeps
    apostrophes inside values from pairing across values: in
    ``{a: 'it's ok', b: 1}`` the quote in ``it's`` is followed by ``s`` so
    it cannot close the string.

    Computed in one reverse pass so each lookup is O(log n) via bisect; a
    forward scan that peeks past whitespace at every quote would be
    quadratic on whitespace-heavy input, and this text is model-emitted.
    """
    closes: list[int] = []
    next_sig = ""  # next non-whitespace char AFTER the current index
    for idx in range(len(s) - 1, -1, -1):
        ch = s[idx]
        if ch == "'" and (next_sig == "" or next_sig in ",}])"):
            closes.append(idx)
        if not ch.isspace():
            next_sig = ch
    closes.reverse()
    return closes


def _scan_gemma4_args_span(
    text: str, open_idx: int, squote_closes: list,
    open_ch: str = "{", close_ch: str = "}",
) -> int:
    """Return the end index (exclusive) of the balanced ``open_ch...close_ch``
    span starting at ``open_idx``, or -1 if no balanced span exists within
    bounds.  Defaults to braces (the canonical ``call:name{...}`` form); the
    parenthesized variant (#1846) passes ``(``/``)`` to find the call's outer
    span — nested ``{}``/``[]`` then pass through as ordinary characters for
    depth purposes (only the tracked pair is counted), while the string-skip
    logic below is character-agnostic so a ``)`` inside any string still
    cannot close the span.

    Iterative single-pass walk that counts brace depth only OUTSIDE string
    literals (``<|"|>``-paired strings, standard JSON double-quoted strings,
    and anchored single-quoted values), so a brace inside string content
    cannot truncate or unbalance the span.
    This deliberately replaces a recursive regex:
    - linear time: recursive alternation patterns degrade quadratically on
      unbalanced model output (measured ~590ms at 80KB), an injection-driven
      CPU burn on a server,
    - iterative: RecursionError is a RuntimeError subclass that no except
      tuple in the parse chain catches, so recursion on deeply nested model
      output would escape as a 500.

    A single-quoted string OPENS only at a value position (previous
    significant char is ``:``, ``,`` or ``[``); a bare apostrophe anywhere
    else (don't, it's) is ordinary content.
    """
    n = len(text)
    depth = 0
    last_sig = ""
    i = open_idx
    limit = min(n, open_idx + _GEMMA4_MAX_ARGS_LEN)
    while i < limit:
        if text.startswith(_GEMMA4_STR_DELIM, i):
            close = text.find(_GEMMA4_STR_DELIM, i + len(_GEMMA4_STR_DELIM))
            if close == -1:
                return -1  # unterminated string: malformed, give up cleanly
            i = close + len(_GEMMA4_STR_DELIM)
            last_sig = '"'
            continue
        ch = text[i]
        if ch == '"':
            # Standard JSON double-quoted string. The <|"|> delimiter is
            # matched by the startswith branch above before we reach here, so
            # a bare ``"`` is an ordinary JSON string open: skip to its closing
            # unescaped quote so a ``}`` inside the value cannot truncate the
            # span (#1854 — without this the suffix remap turned the corrupted
            # parse into an executable call with silently mangled arguments).
            # Honors ``\"`` and ``\\`` so an escaped quote never closes early.
            j = i + 1
            while j < limit:
                if text[j] == "\\":
                    j += 2  # escaped char is literal, never closes the string
                    continue
                if text[j] == '"':
                    break
                j += 1
            else:
                return -1  # unterminated string within bounds: drop cleanly
            i = j + 1
            last_sig = '"'
            continue
        if ch == "'" and last_sig in ":=,[":
            k = bisect.bisect_right(squote_closes, i)
            if k < len(squote_closes):
                i = squote_closes[k] + 1
                last_sig = "'"
                continue
            # No valid close ahead: treat the quote as ordinary content.
        if ch == open_ch:
            depth += 1
        elif ch == close_ch:
            depth -= 1
            if depth == 0:
                return i + 1
        if not ch.isspace():
            last_sig = ch
        i += 1
    return -1


def _gemma4_args_to_json_robust(args_str: str) -> dict:
    """Convert Gemma 4 tool-call args to a Python dict.

    Tries the strict single-pass transcoder first
    (``_gemma4_transcode_to_json``); on failure, falls back to the legacy
    key-anchored recovery (``_gemma4_args_to_json_legacy``) for the one
    case the transcoder deliberately rejects: bare values that themselves
    contain commas, braces, or newlines — long markdown emitted into an
    ``answer:`` argument, observed live on the diffusion lane (#1837).
    The transcoder stops bare values at the first structural separator, so
    only key-anchored capture can recover that shape.
    """
    try:
        return _gemma4_transcode_to_json(args_str)
    except _Gemma4ArgsTooComplexError:
        # Defensive bound breached: reject hard.  The legacy parser ignores
        # these bounds and would parse the input anyway, defeating the DoS
        # guard, so it must NOT see oversized/deeply-nested args.
        raise
    except (ValueError, json.JSONDecodeError):
        # The legacy path's NUL-placeholder forge vector is reintroduced
        # ONLY for ambiguous input the strict transcoder could not parse
        # (e.g. bare multi-comma markdown values, #1837); the common path
        # keeps the transcoder's no-placeholder, injection-safe guarantee.
        return _gemma4_args_to_json_legacy(args_str)


def _gemma4_transcode_to_json(args_str: str) -> dict:
    """Transcode Gemma 4 tool-call args to a dict in a single pass.

    Handles what mlx-lm's parser cannot:
    - bare keys and values (``{location: Tokyo}``)
    - single-quoted values, including commas/colons/braces/apostrophes
      inside them (``{content: 'a, b: c'}``, #1830)
    - ``<|"|>``-delimited strings, arrays, and nested objects

    Implemented as a single-pass transcode to JSON text followed by one
    ``json.loads`` after local length/depth checks.  Every piece of captured
    string content is emitted through ``json.dumps`` and structural characters
    are emitted only by the state machine, so model output cannot inject JSON
    structure.  The legacy
    implementation substituted ``\\x00N\\x00`` placeholders, which literal
    NUL bytes in model output could forge, cross-contaminating argument
    values; transcoding directly leaves nothing to forge.

    Bare values stop at the first ``,``/``}``/``]`` by design: a bare value
    that embeds those characters is ambiguous here, and the caller
    (``_gemma4_args_to_json_robust``) recovers it via the legacy
    key-anchored fallback.
    """
    if len(args_str) > _GEMMA4_MAX_ARGS_LEN:
        raise _Gemma4ArgsTooComplexError("Gemma 4 args too large to parse")

    squote_closes = _squote_close_positions(args_str)
    n = len(args_str)
    out: list[str] = []  # JSON text fragments
    stack: list[str] = []  # open containers: "{" or "["
    expect = "object"  # object | key | value | delim
    i = 0

    def _skip_ws(i: int) -> int:
        while i < n and args_str[i].isspace():
            i += 1
        return i

    def _read_marked_string(i: int):
        """Read a <|"|>- or single-quoted string at i, or return None."""
        if args_str.startswith(_GEMMA4_STR_DELIM, i):
            close = args_str.find(
                _GEMMA4_STR_DELIM, i + len(_GEMMA4_STR_DELIM)
            )
            if close == -1:
                raise ValueError("unterminated Gemma 4 string")
            return (
                args_str[i + len(_GEMMA4_STR_DELIM): close],
                close + len(_GEMMA4_STR_DELIM),
            )
        if args_str[i] == "'":
            k = bisect.bisect_right(squote_closes, i)
            if k < len(squote_closes):
                close = squote_closes[k]
                return args_str[i + 1: close], close + 1
            # No anchored close ahead: not a string, treat as bare content.
        return None

    def _read_json_string(i: int):
        """Read a standard double-quoted JSON string token verbatim."""
        j = i + 1
        while j < n:
            if args_str[j] == "\\":
                j += 2
                continue
            if args_str[j] == '"':
                return args_str[i: j + 1], j + 1
            j += 1
        raise ValueError("unterminated double-quoted string")

    while True:
        i = _skip_ws(i)
        if expect == "object":
            # The outer container is ``{`` for the canonical form and ``(`` for
            # the parenthesized variant (#1846).  Either way it is an object;
            # we always emit ``{`` to JSON and remember the real opener on the
            # stack so its matching closer (``}`` or ``)``) is accepted below.
            if i >= n or args_str[i] not in "{(":
                raise ValueError("Gemma 4 args must start with '{' or '('")
            out.append("{")
            stack.append(args_str[i])
            i += 1
            expect = "key"
        elif expect == "key":
            if i >= n:
                raise ValueError("unterminated object")
            if args_str[i] == _GEMMA4_CLOSE_CHAR[stack[-1]]:
                # Empty object, or tolerated trailing comma.  Closer is ``}``
                # for a ``{`` opener and ``)`` for the paren variant's ``(``.
                if out and out[-1] == ", ":
                    out.pop()
                out.append("}")
                stack.pop()
                i += 1
                expect = "delim"
                continue
            if args_str[i] == '"':
                tok, i = _read_json_string(i)
                key = json.loads(tok)
            else:
                marked = _read_marked_string(i)
                if marked is not None:
                    key, i = marked
                else:
                    # Bare key: everything up to the separator.  ``=`` is a
                    # separator too for the parenthesized kwargs variant
                    # (#1846), so it bounds the key just like ``:``.
                    j = i
                    while j < n and args_str[j] not in ":=,{}[]'\"":
                        j += 1
                    key = args_str[i:j].strip()
                    if not key:
                        raise ValueError("malformed object key")
                    i = j
            i = _skip_ws(i)
            # Accept ``=`` as well as ``:`` — the parenthesized variant (#1846)
            # uses ``key=value``.  This is applied universally (not only at top
            # level), so it does widen the curly grammar to also accept ``=``
            # separators; that is a strict superset, so every valid ``:``-based
            # curly parse is unchanged and a previously-rejected ``{a = 1}`` now
            # succeeds rather than corrupting anything.  ``=`` inside a value is
            # untouched: strings are read atomically and bare values stop only
            # at ``,``/``}``/``]`` (plus ``)`` when a paren container is open).
            if i >= n or args_str[i] not in ":=":
                raise ValueError("expected ':' or '=' after object key")
            out.append(json.dumps(key))
            out.append(": ")
            i += 1
            expect = "value"
        elif expect == "value":
            if i >= n:
                raise ValueError("unterminated value")
            ch = args_str[i]
            if ch == "{" or ch == "[":
                # Depth bound, not recursion: a breach must surface as a
                # clean parse failure on the existing drop path, never as a
                # RecursionError (uncaught by the parse chain's excepts).
                if len(stack) >= _GEMMA4_MAX_DEPTH:
                    raise _Gemma4ArgsTooComplexError(
                        "Gemma 4 args nested too deeply"
                    )
                out.append(ch)
                stack.append(ch)
                i += 1
                expect = "key" if ch == "{" else "value"
                continue
            if ch == "]" and stack and stack[-1] == "[":
                # Empty array, or tolerated trailing comma.
                if out and out[-1] == ", ":
                    out.pop()
                out.append("]")
                stack.pop()
                i += 1
                expect = "delim"
                continue
            if ch == '"':
                tok, i = _read_json_string(i)
                out.append(tok)
                expect = "delim"
                continue
            marked = _read_marked_string(i)
            if marked is not None:
                content, i = marked
                out.append(json.dumps(content))
                expect = "delim"
                continue
            # Bare value: runs to the next structural separator.  When the
            # call uses the parenthesized outer shell (#1846), ``)`` also
            # terminates a bare value so the closing paren of the call is not
            # swallowed (``call:f(units=metric)`` — the value is ``metric``,
            # not ``metric)``).  For the curly form ``)`` stays ordinary
            # content so a value may legitimately contain parentheses
            # (``{expr: f(x)}``).  ``(`` is only ever the outermost container,
            # so its presence on the stack is the reliable signal.
            stops = ",)}]" if "(" in stack else ",}]"
            j = i
            while j < n and args_str[j] not in stops:
                j += 1
            value = args_str[i:j].strip()
            i = j
            if not value:
                raise ValueError("empty value")
            low = value.lower()
            if low in ("true", "false", "null"):
                out.append(low)  # normalize case (models emit True/False)
            else:
                try:
                    json.loads(value)  # already a valid scalar (number, ...)
                    out.append(value)
                except (json.JSONDecodeError, ValueError):
                    out.append(json.dumps(value))
            expect = "delim"
        else:  # expect == "delim"
            if not stack:
                if i < n:
                    raise ValueError("trailing data after args object")
                break
            if i >= n:
                raise ValueError("unterminated args")
            ch = args_str[i]
            if ch == ",":
                out.append(", ")
                i += 1
                # After a comma, an object (``{`` or the paren outer ``(``)
                # expects a key; an array expects a value.
                expect = "key" if stack[-1] in "{(" else "value"
            elif stack[-1] in "{(" and ch == _GEMMA4_CLOSE_CHAR[stack[-1]]:
                # Object close: ``}`` for ``{``, ``)`` for the paren outer.
                out.append("}")
                stack.pop()
                i += 1
            elif ch == "]" and stack[-1] == "[":
                out.append("]")
                stack.pop()
                i += 1
            else:
                raise ValueError("malformed args structure")

    result = json.loads("".join(out))
    if not isinstance(result, dict):
        raise ValueError("Gemma 4 args did not parse to an object")
    return result


def _gemma4_args_to_json_legacy(args_str: str) -> dict:
    """Legacy regex-based Gemma 4 args parser (upstream #1837).

    Kept as the last-resort fallback behind ``_gemma4_transcode_to_json``.
    Its value over the transcoder is step 6: key-anchored value capture for
    bare values that themselves contain commas, braces, or newlines (long
    markdown emitted into an ``answer:`` argument, observed live on the
    diffusion lane).  The transcoder stops bare values at the first
    separator, so this is the only path that recovers that shape.

    Carries the placeholder mechanism (``\\x00N\\x00``) the transcoder was
    written to avoid; it runs only on input the transcoder already rejected.
    """
    import regex

    # 1. Extract <|"|>-delimited strings and replace with placeholders
    strings: list[str] = []

    def _capture(m):
        strings.append(m.group(1))
        return f"\x00{len(strings) - 1}\x00"

    text = regex.sub(r'<\|"\|>(.*?)<\|"\|>', _capture, args_str, flags=regex.DOTALL)

    # 2. Quote bare keys (allow whitespace after { or ,)
    text = regex.sub(r"(?<=[{,])\s*(\w+)\s*:", r' "\1":', text)

    # 3. Restore captured strings as properly escaped JSON strings
    for i, s in enumerate(strings):
        text = text.replace(f"\x00{i}\x00", json.dumps(s))

    # 4. Try json.loads — works when all values are already valid JSON primitives
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 5. Quote bare string values that are not numbers, booleans, or null
    def _quote_bare(m):
        value = m.group(2).strip()
        suffix = m.group(3)
        if value.lower() in ("true", "false", "null"):
            return f": {value}{suffix}"
        try:
            json.loads(value)
            return f": {value}{suffix}"
        except (json.JSONDecodeError, ValueError):
            return f": {json.dumps(value)}{suffix}"

    # Keep the pre-step-5 text: if step 5 fails, its partial quoting has
    # corrupted multi-line bare values and step 6 must start clean.
    pre_quote_text = text
    text = regex.sub(
        r"(:\s*)([^\",\[\]{}\s][^,}]*?)(\s*[,}])", _quote_bare, text
    )
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 6. Last resort: key-anchored value capture. Bare values that
    # themselves contain commas, braces, or newlines (e.g. long markdown
    # emitted into an ``answer:`` argument — observed live on the
    # diffusion lane) defeat the per-pair regex in step 5. Anchor on the
    # quoted keys produced by step 2 and treat everything between a
    # key's colon and the next key (or the end) as that key's value.
    # Operates on the pre-step-5 text so step 5's partial quoting cannot
    # corrupt the captured values.
    inner = pre_quote_text.strip()
    if inner.startswith("{") and inner.endswith("}"):
        inner = inner[1:-1]
    key_pat = regex.compile(r'"([A-Za-z_]\w*)"\s*:')
    key_matches = list(key_pat.finditer(inner))
    if not key_matches:
        return json.loads(text)  # re-raise original-style error
    result: dict = {}
    for i, km in enumerate(key_matches):
        value_start = km.end()
        value_end = (
            key_matches[i + 1].start() if i + 1 < len(key_matches) else len(inner)
        )
        raw_value = inner[value_start:value_end].strip()
        if i + 1 < len(key_matches):
            raw_value = raw_value.rstrip().rstrip(",").rstrip()
        try:
            result[km.group(1)] = json.loads(raw_value)
        except (json.JSONDecodeError, ValueError):
            result[km.group(1)] = raw_value
    return result


def _parse_gemma4_tool_call_fallback(text: str) -> Union[dict, list]:
    """Robust fallback parser for Gemma 4 ``call:name{args}`` format.

    Activated only for Gemma 4 models (guarded by the ``tool_call_start``
    check at the call site).  Extends mlx-lm's parser to handle:
    - colons / dots / hyphens in function names (namespaced MCP tools,
      e.g. ``call:google:mcp:text_generation:create-pdf-file``, #1830)
    - bare string values without ``<|"|>`` delimiters
    - single-quoted values, including commas, colons, braces and
      apostrophes inside them (#1830)
    - the parenthesized kwargs variant ``call:name(key=value, ...)`` that
      Gemma 4 26B degrades to under instruction-dense agentic load (#1846).
      The nested grammar is identical to the curly form, so only the outer
      ``()`` shell and the ``=`` separator are new; both head forms are
      collected and processed in document order (see below).
    - degenerate ``call:`` prefixes from the diffusion lane's parallel
      denoising, which can drop a token from the opening (observed live:
      ``calldone{...}`` — missing colon — and ``:done{...}`` — missing
      ``call``, #1837).  ``_GEMMA4_CALL_HEAD`` matches these; the text is
      already marker-delimited (between ``<|tool_call>`` and
      ``<tool_call|>``), so the permissive prefix cannot misfire on prose.

    Name remapping onto registered tools is deliberately NOT done here:
    that is a post-parse concern handled by ``_remap_tool_call_names`` so
    it covers every producer path (native parser, this fallback, XML
    recovery, thinking-content promotion), not just this one.
    """
    squote_closes = _squote_close_positions(text)

    # Collect heads from both the canonical curly form and the parenthesized
    # variant (#1846), then process them in document order so the consumed-span
    # dedup below holds across BOTH forms: a curly head that falls inside an
    # already-consumed paren span (the inner ``{...}`` of ``call:f(a={...})``)
    # is string/structure content, not a sibling call, and must be skipped.
    heads = [(m.start(), m, "{", "}") for m in _GEMMA4_CALL_HEAD.finditer(text)]
    heads += [
        (m.start(), m, "(", ")")
        for m in _GEMMA4_CALL_HEAD_PAREN.finditer(text)
    ]
    heads.sort(key=lambda h: h[0])

    results = []
    consumed_until = 0
    for start, m, open_ch, close_ch in heads:
        # A head inside an already-consumed args span is string content
        # (e.g. quoted prose mentioning a tool call), not a sibling call.
        if start < consumed_until:
            continue
        open_idx = m.end() - 1
        end = _scan_gemma4_args_span(
            text, open_idx, squote_closes, open_ch, close_ch
        )
        if end == -1:
            continue
        args_str = text[open_idx:end]
        try:
            arguments = _gemma4_args_to_json_robust(args_str)
        except (ValueError, json.JSONDecodeError, RecursionError):
            continue  # one malformed call must not drop its siblings
        if not isinstance(arguments, dict):
            continue
        results.append({"name": m.group(1), "arguments": arguments})
        consumed_until = end

    if not results:
        raise ValueError("No function call found in Gemma 4 format")
    return results[0] if len(results) == 1 else results


def _remap_tool_call_names(
    tool_calls: List[ToolCall], tools: Optional[List]
) -> None:
    """Remap namespace-prefixed emitted tool names onto registered tools.

    Gemma 4 emits names like ``google:mcp:text_generation:create-pdf-file``
    for a tool registered as ``create-pdf-file`` (#1830); clients match by
    exact name, so the call would be unusable.  Runs post-parse so every
    producer path is covered and so the behavior survives changes to
    mlx-lm's native parser (which currently rejects colon names and routes
    these to the fallback, but may not forever).

    Rule: remap only when the emitted name matches no registered tool AND
    exactly one registered tool is a ``:``-boundary suffix of it; on zero
    or several candidates keep the name verbatim.  The comparison is
    boundary-aligned by construction (split on ':'), never str.endswith:
    a bare endswith would let a crafted emission like 'evilcreate-pdf-file'
    coerce into a registered 'create-pdf-file' (model output is
    attacker-influenceable via prompt injection).
    """
    if not tool_calls or not tools:
        return
    valid_names = _extract_tool_names(tools)
    if not valid_names:
        return
    for tc in tool_calls:
        name = tc.function.name if tc.function else ""
        if not name or name in valid_names or ":" not in name:
            continue
        parts = name.split(":")
        suffixes = {":".join(parts[i:]) for i in range(1, len(parts))}
        candidates = suffixes & valid_names
        if len(candidates) == 1:
            target = next(iter(candidates))
            logger.info(
                "Remapped namespaced tool call name %r to registered "
                "tool %r",
                name[:200],
                target,
            )
            tc.function.name = target


def parse_tool_calls(
    text: str,
    tokenizer: Any,
    tools: Optional[List] = None,
) -> Tuple[str, Optional[List[ToolCall]]]:
    """
    Parse tool calls from model output.

    Uses mlx-lm's TokenizerWrapper tool parser if available, otherwise
    falls back to generic XML tool call parsing for models like GLM.

    Emitted names that match no registered tool are conservatively remapped
    onto registered tools afterwards (see _remap_tool_call_names); doing it
    here, at the single post-parse chokepoint, covers every producer path
    including the thinking-content promotion in
    extract_tool_calls_with_thinking, whose exact-name validity filter would
    otherwise silently drop cleanly-parsed namespaced calls (#1830).

    Args:
        text: Raw model output text
        tokenizer: mlx-lm's TokenizerWrapper (required)
        tools: Tool definitions for type conversion (optional)

    Returns:
        Tuple of (cleaned_text, tool_calls or None)
        - cleaned_text: Text with tool call tags and thinking tags removed
        - tool_calls: List of ToolCall objects, or None if no tool calls found
    """
    cleaned_text, tool_calls = _parse_tool_calls_impl(text, tokenizer, tools)
    if tool_calls:
        _remap_tool_call_names(tool_calls, tools)
    return cleaned_text, tool_calls


def _parse_tool_calls_impl(
    text: str,
    tokenizer: Any,
    tools: Optional[List] = None,
) -> Tuple[str, Optional[List[ToolCall]]]:
    """parse_tool_calls body, pre-remap. See the public wrapper's docstring."""
    cleaned_text = text

    # Remove thinking tags if present (reasoning models)
    cleaned_text = re.sub(
        r"<think>.*?</think>", "", cleaned_text, flags=re.DOTALL
    ).strip()

    # Try mlx-lm's native tool parser first
    if getattr(tokenizer, "has_tool_calling", False):
        tool_call_start = tokenizer.tool_call_start
        tool_call_end = tokenizer.tool_call_end
        tool_parser = tokenizer.tool_parser

        if tool_call_start is not None and tool_parser is not None:
            tool_calls = []
            start_escaped = re.escape(tool_call_start)

            if tool_call_end:
                # Paired markers (e.g. <tool_call>...</tool_call>)
                end_escaped = re.escape(tool_call_end)
                pattern = rf"{start_escaped}(.*?){end_escaped}"
                matches = re.findall(pattern, text, re.DOTALL)
            else:
                # One-sided marker (e.g. Mistral/Devstral "[TOOL_CALLS]"):
                # split on the start marker and parse each segment.
                # The model emits: [TOOL_CALLS]name[ARGS]{...}[TOOL_CALLS]name2[ARGS]{...}
                parts = re.split(start_escaped, text)
                # First part is pre-marker text, rest are tool call segments
                matches = [p for p in parts[1:] if p.strip()]

            for match in matches:
                try:
                    parsed = tool_parser(match.strip(), tools)
                    # MiniMax M2 parser returns a list when a single
                    # <minimax:tool_call> block contains multiple <invoke>s.
                    items = parsed if isinstance(parsed, list) else [parsed]
                    for p in items:
                        name = p.get("name", "")
                        arguments = p.get("arguments", {})
                        tool_calls.append(
                            ToolCall(
                                id=f"call_{uuid.uuid4().hex[:8]}",
                                type="function",
                                function=FunctionCall(
                                    name=name,
                                    arguments=_serialize_tool_call_arguments(arguments),
                                ),
                            )
                        )
                except (
                    ValueError,
                    json.JSONDecodeError,
                    AttributeError,
                    KeyError,
                    SyntaxError,
                    TypeError,
                ) as primary_err:
                    # Gemma 4 only: try robust fallback that handles bare
                    # string values and colons in function names.
                    gemma4_handled = False
                    if tool_call_start == "<|tool_call>":
                        try:
                            parsed = _parse_gemma4_tool_call_fallback(
                                match.strip()
                            )
                            items = (
                                parsed if isinstance(parsed, list) else [parsed]
                            )
                            for p in items:
                                name = p.get("name", "")
                                arguments = p.get("arguments", {})
                                tool_calls.append(
                                    ToolCall(
                                        id=f"call_{uuid.uuid4().hex[:8]}",
                                        type="function",
                                        function=FunctionCall(
                                            name=name,
                                            arguments=_serialize_tool_call_arguments(
                                                arguments
                                            ),
                                        ),
                                    )
                                )
                            gemma4_handled = True
                        except (
                            ValueError,
                            json.JSONDecodeError,
                            KeyError,
                            SyntaxError,
                            TypeError,
                        ):
                            pass

                    if gemma4_handled:
                        continue

                    # Per-match XML fallback: regex-only, no ast.literal_eval,
                    # recovers Qwen/GLM/Hermes-JSON formats. Prevents silent
                    # drop when the native parser raises (e.g. ast.literal_eval
                    # SyntaxError on non-Python-literal parameter values).
                    fb_wrapped = f"<tool_call>{match}</tool_call>"
                    _, fb_calls = _parse_xml_tool_calls(fb_wrapped, tools)
                    if fb_calls:
                        tool_calls.extend(fb_calls)
                        logger.warning(
                            "Native tool parser failed (%s: %s), "
                            "recovered via XML fallback. Match: %r",
                            type(primary_err).__name__,
                            primary_err,
                            match[:200],
                        )
                    else:
                        logger.warning(
                            "Native tool parser failed (%s: %s) and XML "
                            "fallback could not recover. Dropping match: %r",
                            type(primary_err).__name__,
                            primary_err,
                            match[:200],
                        )
                    continue

            if tool_calls:
                if tool_call_end:
                    cleaned_text = re.sub(
                        rf"{start_escaped}.*?{re.escape(tool_call_end)}",
                        "",
                        cleaned_text,
                        flags=re.DOTALL,
                    ).strip()
                else:
                    # One-sided: everything from first marker to end is tool calls
                    idx = cleaned_text.find(tool_call_start)
                    if idx >= 0:
                        cleaned_text = cleaned_text[:idx].strip()
                return cleaned_text, tool_calls

    # Fallback: parse XML <tool_call> tags (GLM, Qwen, generic formats)
    if "<tool_call>" in cleaned_text:
        return _parse_xml_tool_calls(cleaned_text, tools)

    # Fallback: namespaced tool_call tags (e.g. <minimax:tool_call>)
    ns_match = re.search(r"<([A-Za-z_][\w.-]*):tool_call>", cleaned_text)
    if ns_match:
        ns = ns_match.group(1)
        return _parse_namespaced_tool_calls(cleaned_text, ns, tools)

    # Fallback: Hermes-style tool calls (<|tool_call_start|>[func(args)]<|tool_call_end|>)
    if "<|tool_call_start|>" in cleaned_text:
        hermes_result = _parse_hermes_tool_calls(cleaned_text)
        if hermes_result[1] is not None:
            return hermes_result

    # Fallback: bracket tool call formats (from text-formatted history)
    if "[Calling tool:" in cleaned_text or "[Tool call:" in cleaned_text:
        return _parse_bracket_tool_calls(cleaned_text)

    # All parsing attempts exhausted. Strip known tool-call markers so raw
    # control markup never leaks into the API response.  Models whose markers
    # overlap with the generic ``<tool_call>`` tag already returned above via
    # Branch 2 (_parse_xml_tool_calls), so this only affects models with
    # unique markers (Gemma 4, Mistral, Pythonic, Kimi K2, Longcat, etc.).
    if getattr(tokenizer, "has_tool_calling", False):
        _start = getattr(tokenizer, "tool_call_start", None)
        _end = getattr(tokenizer, "tool_call_end", None)
        if _start and _end:
            s_esc = re.escape(_start)
            e_esc = re.escape(_end)
            stripped = re.findall(
                rf"{s_esc}(.*?){e_esc}", cleaned_text, flags=re.DOTALL
            )
            if stripped:
                logger.warning(
                    "Tool call markers found but parsing failed, "
                    "stripping markers. Raw content: %s",
                    stripped,
                )
            cleaned_text = re.sub(
                rf"{s_esc}.*?{e_esc}", "", cleaned_text, flags=re.DOTALL
            ).strip()
        elif _start:
            idx = cleaned_text.find(_start)
            if idx >= 0:
                logger.warning(
                    "Tool call start marker found but parsing failed, "
                    "stripping marker. Raw content: %s",
                    cleaned_text[idx:],
                )
                cleaned_text = cleaned_text[:idx].strip()

    # Strip Hermes markers if still present (models without has_tool_calling)
    if "<|tool_call_start|>" in cleaned_text:
        cleaned_text = re.sub(
            r"<\|tool_call_start\|>.*?<\|tool_call_end\|>",
            "",
            cleaned_text,
            flags=re.DOTALL,
        ).strip()

    return cleaned_text, None


def sanitize_tool_call_markup(text: str, tokenizer: Any) -> str:
    """Remove tool-call control markup while preserving surrounding prose."""
    if not text:
        return ""

    stream_filter = ToolCallStreamFilter(tokenizer)
    cleaned = stream_filter.feed(text)
    cleaned += stream_filter.finish()
    return cleaned.strip()


def _extract_tool_names(tools: List) -> set:
    """Extract function names from OpenAI-format tool definitions."""
    names = set()
    for tool in tools:
        if isinstance(tool, dict):
            func = tool.get("function", {})
            if isinstance(func, dict):
                name = func.get("name")
                if name:
                    names.add(name)
    return names


def extract_tool_calls_with_thinking(
    thinking_content: str,
    regular_content: str,
    tokenizer: Any,
    tools: Optional[List] = None,
) -> ToolCallExtraction:
    """Extract tool calls while keeping a sanitized reasoning transcript.

    When tool calls are found in thinking content (not regular content),
    the ``tools`` parameter controls validation:

    * ``None`` (default) — no tools list was provided.  Thinking-embedded
      calls are kept only when ``regular_content`` is empty (the model
      produced no competing prose).  Otherwise they are dropped as
      potential hallucinated reasoning.
    * ``[]`` — "no tools allowed".  All thinking-embedded calls are
      dropped regardless of ``regular_content``.
    * Non-empty list — name matching is the sole discriminator.
      Calls whose name matches a provided tool are promoted regardless
      of whether regular text was also produced.
    """
    cleaned_text, tool_calls = parse_tool_calls(regular_content, tokenizer, tools)
    cleaned_thinking = sanitize_tool_call_markup(thinking_content, tokenizer)
    tool_calls_from_thinking = False

    if not tool_calls and thinking_content:
        _, tool_calls = parse_tool_calls(thinking_content, tokenizer, tools)
        tool_calls_from_thinking = bool(tool_calls)

        # Guard: validate thinking-embedded tool calls.
        #
        # Three cases:
        # 1. tools is None (not provided) AND regular text exists → drop.
        #    The call is unvalidated and could be hallucinated reasoning.
        # 2. tools is None AND no regular text → keep.  The model clearly
        #    intended a tool invocation (no competing prose).
        # 3. tools is a list (including empty) → name matching is the sole
        #    discriminator.  An empty list means "no tools allowed" so all
        #    calls are dropped.  A non-empty list filters by name, regardless
        #    of whether regular text was also produced.  The previous "regular
        #    text means just reasoning" heuristic was wrong for models
        #    (Qwen3-Coder) that genuinely place tool calls in thinking.
        # See https://github.com/jundot/omlx/issues/1392
        if tool_calls:
            if tools is None:
                if regular_content.strip():
                    tool_calls = None
                    tool_calls_from_thinking = False
            else:
                valid_names = _extract_tool_names(tools)
                tool_calls = [tc for tc in tool_calls if tc.function.name in valid_names]
                if not tool_calls:
                    tool_calls = None
                    tool_calls_from_thinking = False

    return ToolCallExtraction(
        cleaned_text=cleaned_text,
        tool_calls=tool_calls,
        cleaned_thinking=cleaned_thinking,
        tool_calls_from_thinking=tool_calls_from_thinking,
    )


def parse_tool_calls_with_thinking_fallback(
    thinking_content: str,
    regular_content: str,
    tokenizer: Any,
    tools: Optional[List] = None,
) -> Tuple[str, Optional[List[ToolCall]]]:
    """Parse tool calls from content, falling back to thinking if none found.

    Small reasoning models sometimes generate tool call XML inside <think>
    blocks instead of after </think>. This function first tries the normal
    content, then falls back to parsing from thinking content.

    Args:
        thinking_content: Text extracted from <think>...</think> blocks.
        regular_content: Text outside thinking blocks.
        tokenizer: mlx-lm's TokenizerWrapper.
        tools: Tool definitions for type conversion (optional).

    Returns:
        Tuple of (cleaned_text, tool_calls or None).
        cleaned_text comes from regular_content only (thinking text is
        never promoted to content).
    """
    result = extract_tool_calls_with_thinking(
        thinking_content,
        regular_content,
        tokenizer,
        tools,
    )
    return result.cleaned_text, result.tool_calls


class ToolCallStreamFilter:
    """Streaming filter that suppresses tool-call markup from content deltas.

    Detects known tool-call start envelopes during streaming and suppresses
    control markup from assistant-visible content. Supports tokenizer-defined
    delimiters, namespaced XML envelopes, and high-confidence bracket-format
    envelopes handled by ``parse_tool_calls``.

    Suppression is envelope-bounded: control markup is removed, then visible
    prose after a closed envelope continues streaming normally.

    Args:
        tokenizer: The model's tokenizer. Uses tokenizer-defined
            ``tool_call_start`` when available.
    """

    def __init__(self, tokenizer: Any):
        marker = getattr(tokenizer, "tool_call_start", None)
        marker_end = getattr(tokenizer, "tool_call_end", None)
        # Normalize None-like values but preserve empty strings.
        if marker is None:
            marker = ""
        if marker_end is None:
            marker_end = ""
        self._marker_pairs: List[Tuple[str, str]] = [
            ("]<]minimax[>[<tool_call>", "]<]minimax[>[</tool_call>"),
            ("<|tool_call_start|>", "<|tool_call_end|>"),
            ("<tool_call>", "</tool_call>"),
        ]
        self._suppress_after_markers: List[str] = []
        if marker:
            if marker_end:
                self._marker_pairs.insert(0, (marker, marker_end))
            else:
                # One-sided markers (e.g. Mistral "[TOOL_CALLS]" with no
                # end marker): suppress everything after the start marker.
                self._suppress_after_markers.append(marker)
        # Gemma 4 can emit a bare close token outside a matched tool-call
        # envelope. Do not apply this to XML-style closers like </tool_call>,
        # which may appear as literal prose.
        is_gemma4_tool_marker = (
            marker == "<|tool_call>" and marker_end == "<tool_call|>"
        )
        self._stray_close_markers: List[str] = (
            [marker_end] if is_gemma4_tool_marker else []
        )
        self._orphan_close_markers: List[str] = ["<|tool_call_end|>"]
        if marker_end and not self._is_xml_close_marker(marker_end):
            self._orphan_close_markers.append(marker_end)
        self._orphan_close_markers = list(dict.fromkeys(self._orphan_close_markers))
        self._namespaced_open_re = re.compile(r"<([A-Za-z_][\w.-]*):tool_call>")
        self._bracket_prefixes = ["[Calling tool:", "[Tool call:"]
        self._bracket_call_re = re.compile(
            r"^\[(?:Calling tool|Tool call):\s*([A-Za-z_][\w.-]*)(?:\(({.*?})\))?\]",
            re.DOTALL,
        )
        self._buffer = ""
        self._suppressing_until: Optional[str] = None
        self._suppressing = False
        self._pending_envelope_parts: List[str] = []
        self._pending_start_marker: Optional[str] = None
        self._recovery_candidate = ""

    @staticmethod
    def _is_xml_close_marker(marker: str) -> bool:
        return marker.startswith("</") and marker.endswith(">")

    @property
    def active(self) -> bool:
        """Whether this filter should run for tool-enabled streams."""
        return True

    def take_recovery_candidate(self) -> str:
        """Return and clear an unterminated paired envelope captured at EOF.

        The caller must only surface this text after final tool parsing confirms
        that no structured tool call was recovered. This keeps valid tool calls
        hidden even when another parser can recover malformed outer markup.
        """
        candidate = self._recovery_candidate
        self._recovery_candidate = ""
        return candidate

    def _clear_pending_envelope(self) -> None:
        self._pending_envelope_parts = []
        self._pending_start_marker = None

    def _find_start_envelope(
        self, text: str
    ) -> Optional[Tuple[int, int, Optional[str]]]:
        """Find earliest complete opening envelope.

        Returns:
            tuple(index, consume_len, close_marker_or_none)
            - close_marker_or_none is a close marker to wait for, or ``None``
              when the whole envelope is already contained in consume_len.
        """
        starts: List[Tuple[int, int, Optional[str]]] = []

        for marker, close in self._marker_pairs:
            idx = text.find(marker)
            if idx >= 0:
                starts.append((idx, len(marker), close))

        for close in self._orphan_close_markers:
            close_idx = text.find(close)
            if close_idx >= 0:
                starts.append((close_idx, len(close), None))

        ns_match = self._namespaced_open_re.search(text)
        if ns_match:
            ns = ns_match.group(1)
            starts.append(
                (ns_match.start(), len(ns_match.group(0)), f"</{ns}:tool_call>")
            )

        for bp in self._bracket_prefixes:
            bracket_idx = text.find(bp)
            while bracket_idx >= 0:
                bracket_candidate = text[bracket_idx:]
                bracket_match = self._bracket_call_re.match(bracket_candidate)
                if bracket_match:
                    starts.append((bracket_idx, bracket_match.end(), None))
                bracket_idx = text.find(bp, bracket_idx + 1)

        # One-sided markers: suppress from start marker to end of buffer.
        for sa_marker in self._suppress_after_markers:
            idx = text.find(sa_marker)
            if idx >= 0:
                starts.append((idx, len(text) - idx, "__suppress_permanently__"))

        if not starts:
            return None
        return min(starts, key=lambda x: x[0])

    @staticmethod
    def _partial_prefix_len(text: str, marker: str) -> int:
        """Longest suffix of text that is a proper prefix of marker."""
        max_len = min(len(text), len(marker) - 1)
        for n in range(max_len, 0, -1):
            if text.endswith(marker[:n]):
                return n
        return 0

    @staticmethod
    def _could_be_partial_namespaced_open(candidate: str) -> bool:
        """Return True if candidate could prefix a namespaced <ns:tool_call> tag."""
        if not candidate.startswith("<"):
            return False
        if ">" in candidate:
            return False

        body = candidate[1:]
        if not body:
            return True
        if body.startswith("/"):
            return False

        if ":" not in body:
            return re.match(r"^[A-Za-z_][\w.-]*$", body) is not None

        ns, suffix = body.split(":", 1)
        if not re.match(r"^[A-Za-z_][\w.-]*$", ns):
            return False
        return "tool_call".startswith(suffix)

    def _partial_suffix_len(self, text: str) -> int:
        """Length of trailing suffix that might be an opening-marker prefix."""
        keep = 0
        for marker, _close in self._marker_pairs:
            keep = max(keep, self._partial_prefix_len(text, marker))

        last_lt = text.rfind("<")
        if last_lt >= 0:
            candidate = text[last_lt:]
            if self._could_be_partial_namespaced_open(candidate):
                keep = max(keep, len(candidate))

        # Partial prefix detection for bracket markers (e.g. "[", "[C",
        # "[Cal" could be start of "[Calling tool:" or "[Tool call:").
        for bp in self._bracket_prefixes:
            keep = max(keep, self._partial_prefix_len(text, bp))
        # Same for suppress-after markers (e.g. "[TOOL" for "[TOOL_CALLS]").
        for sa_marker in self._suppress_after_markers:
            keep = max(keep, self._partial_prefix_len(text, sa_marker))
        # Hold partial prefix of a stray-close marker so it reassembles before
        # the strip check — prevents the "hello<tool_call|" + ">" split leak.
        for close_marker in self._orphan_close_markers:
            keep = max(keep, self._partial_prefix_len(text, close_marker))

        bracket_idx = -1
        for bp in self._bracket_prefixes:
            idx = text.rfind(bp)
            if idx > bracket_idx:
                bracket_idx = idx
        if bracket_idx >= 0:
            bracket_candidate = text[bracket_idx:]
            # Hold unresolved bracket prefix until we can classify parseable
            # envelope vs literal prose.
            if "]" not in bracket_candidate:
                keep = max(keep, len(bracket_candidate))
                # Do not cap unresolved bracket candidates: capping can leak
                # raw control markup once the prefix grows past the cap.
                return keep

        # Cap retained suffix window to avoid unbounded buffering on malformed text.
        return min(keep, 128)

    def _should_drop_tail_at_finish(self, tail: str) -> bool:
        """Whether unresolved tail should be suppressed under strict mode."""
        if not tail:
            return False

        for marker, _close in self._marker_pairs:
            if marker.startswith(tail):
                # MiniMax M3 markers start with ``]``. A single closing
                # bracket at end-of-stream is much more likely to be literal
                # prose than an incomplete MiniMax control marker.
                if tail == "]":
                    continue
                return True

        for close_marker in self._orphan_close_markers:
            if close_marker.startswith(tail):
                return True

        # Drop unresolved bracket tool-call prefixes
        for bp in self._bracket_prefixes:
            if tail.startswith(bp):
                return True

        # Drop unresolved suppress-after marker prefixes
        for sa_marker in self._suppress_after_markers:
            if sa_marker.startswith(tail) or tail.startswith(sa_marker):
                return True

        if not tail.startswith("<"):
            return False
        if ">" in tail:
            return False

        body = tail[1:]
        if not body:
            return True
        if body.startswith("/"):
            return False

        if ":" not in body:
            # Preserve plain literal tails like "<alpha".
            return False

        ns, suffix = body.split(":", 1)
        if not re.match(r"^[A-Za-z_][\w.-]*$", ns):
            return False
        return "tool_call".startswith(suffix)

    def _sanitize_prefix_before_suppression(self, text: str) -> str:
        """Strip unresolved bracket-control prefixes while preserving prose."""
        if not any(bp in text for bp in self._bracket_prefixes):
            return text

        out: List[str] = []
        cursor = 0
        while cursor < len(text):
            bracket_idx = -1
            bracket_prefix = ""
            for bp in self._bracket_prefixes:
                idx = text.find(bp, cursor)
                if idx >= 0 and (bracket_idx < 0 or idx < bracket_idx):
                    bracket_idx = idx
                    bracket_prefix = bp
            if bracket_idx < 0:
                out.append(text[cursor:])
                break

            out.append(text[cursor:bracket_idx])
            after_prefix = bracket_idx + len(bracket_prefix)
            close_idx = text.find("]", after_prefix)
            if close_idx < 0:
                # Drop only the marker token; keep following prose.
                cursor = after_prefix
                continue

            # Preserve balanced literal bracket text that is not being suppressed.
            out.append(text[bracket_idx : close_idx + 1])
            cursor = close_idx + 1

        return "".join(out)

    def feed(self, text: str) -> str:
        """Feed a content delta, return the portion safe to emit."""
        if self._suppressing or not text:
            return ""
        if not self.active:
            return text

        self._buffer += text
        out: List[str] = []

        while self._buffer:
            if self._suppressing_until == "__suppress_permanently__":
                self._suppressing = True
                self._suppressing_until = None
                self._buffer = ""
                break

            if self._suppressing_until is not None:
                end_idx = self._buffer.find(self._suppressing_until)
                if end_idx < 0:
                    keep = self._partial_prefix_len(
                        self._buffer, self._suppressing_until
                    )
                    if keep:
                        self._pending_envelope_parts.append(self._buffer[:-keep])
                        self._buffer = self._buffer[-keep:]
                    else:
                        self._pending_envelope_parts.append(self._buffer)
                        self._buffer = ""
                    break
                self._buffer = self._buffer[end_idx + len(self._suppressing_until) :]
                self._suppressing_until = None
                self._clear_pending_envelope()
                continue

            start = self._find_start_envelope(self._buffer)
            if start:
                idx, consume_len, close_marker = start
                opening_marker = self._buffer[idx : idx + consume_len]
                if idx > 0:
                    out.append(
                        self._sanitize_prefix_before_suppression(self._buffer[:idx])
                    )
                self._buffer = self._buffer[idx + consume_len :]
                if close_marker is not None:
                    self._suppressing_until = close_marker
                    if close_marker != "__suppress_permanently__":
                        # Recover the exact opening bytes, including dynamic
                        # namespace markers, if the matching close never arrives.
                        self._pending_envelope_parts = [opening_marker]
                        self._pending_start_marker = opening_marker
                continue

            keep = self._partial_suffix_len(self._buffer)
            if keep == 0:
                out.append(self._buffer)
                self._buffer = ""
                break
            if len(self._buffer) > keep:
                out.append(self._buffer[:-keep])
                self._buffer = self._buffer[-keep:]
            break

        result = "".join(out)
        for close in self._stray_close_markers:
            if close in result:
                result = result.replace(close, "")
        return result

    def finish(self) -> str:
        """Flush remaining safe buffer content.

        In clean-output strict mode, unresolved marker-like suffixes are dropped
        so partial control markup does not leak into user-visible text.
        """
        if self._suppressing:
            self._buffer = ""
            self._suppressing_until = None
            self._clear_pending_envelope()
            return ""

        if self._suppressing_until is not None:
            self._pending_envelope_parts.append(self._buffer)
            candidate = "".join(self._pending_envelope_parts)
            start_marker = self._pending_start_marker or "<unknown>"
            self._buffer = ""
            self._suppressing_until = None
            self._clear_pending_envelope()
            if candidate:
                self._recovery_candidate = candidate
                logger.warning(
                    "Unclosed tool-call envelope at end of stream; "
                    "withheld %d characters are available for content recovery "
                    "(start_marker=%.80r)",
                    len(candidate),
                    start_marker,
                )
            return ""

        keep = self._partial_suffix_len(self._buffer)
        if keep >= len(self._buffer):
            tail = self._buffer
            self._buffer = ""
            if self._should_drop_tail_at_finish(tail):
                return ""
            return tail

        if keep:
            buf = self._buffer[:-keep]
            tail = self._buffer[-keep:]
            if not self._should_drop_tail_at_finish(tail):
                buf += tail
        else:
            buf = self._buffer
        self._buffer = ""
        for close in self._stray_close_markers:
            if close in buf:
                buf = buf.replace(close, "")
        return buf


def convert_tools_for_template(tools: Optional[List]) -> Optional[List[dict]]:
    """
    Convert OpenAI tools format to format expected by tokenizer.apply_chat_template.

    OpenAI format:
    [{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}]

    Template format (commonly used by models):
    [{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}]

    Args:
        tools: List of ToolDefinition objects or dicts in OpenAI format

    Returns:
        List of tool definitions in template format, or None if no tools
    """
    if not tools:
        return None

    converted = []
    for tool in tools:
        # Handle both Pydantic models and dicts
        if isinstance(tool, dict):
            tool_type = tool.get("type")
            tool_func = tool.get("function")
        else:
            tool_type = getattr(tool, "type", None)
            tool_func = getattr(tool, "function", None)

        if tool_type == "function" and tool_func:
            # Handle function as dict or Pydantic model
            if isinstance(tool_func, dict):
                func_name = tool_func.get("name", "")
                func_desc = tool_func.get("description", "")
                func_params = tool_func.get(
                    "parameters", {"type": "object", "properties": {}}
                )
            else:
                func_name = getattr(tool_func, "name", "")
                func_desc = getattr(tool_func, "description", "")
                func_params = getattr(
                    tool_func, "parameters", {"type": "object", "properties": {}}
                )

            if func_params is None:
                func_params = {"type": "object", "properties": {}}

            converted.append(
                {
                    "type": "function",
                    "function": {
                        "name": func_name,
                        "description": _template_safe_description(func_desc),
                        "parameters": _copy_schema_with_template_defaults(
                            func_params, is_schema=False
                        ),
                    },
                }
            )

    return converted if converted else None


# Parameter names that collide with JSON Schema keywords.
# Gemma 4 confuses these with schema-level fields and drops them from
# tool call output.  We rename them before the chat template and restore
# them after parsing the model's response.
_GEMMA4_COLLIDING_PARAMS = {"description"}
_GEMMA4_RENAME_PREFIX = "param_"


def enrich_tool_params_for_gemma4(tools: list[dict]) -> list[dict]:
    """Fix tool schemas for Gemma 4 models.

    1. Renames parameters whose names collide with JSON Schema keywords
       (e.g. ``description`` -> ``param_description``) so Gemma 4 doesn't
       confuse them with schema-level fields.
    2. Adds explicit descriptions to required parameters that lack them.

    Use :func:`restore_gemma4_param_names` on tool call arguments to
    reverse the renaming before returning them to the caller.
    """
    enriched = []
    for tool in tools:
        tool = dict(tool)
        func = dict(tool.get("function", {}))
        params = func.get("parameters", {})
        if isinstance(params, dict) and "properties" in params:
            params = dict(params)
            old_props = params.get("properties", {})
            required = list(params.get("required", []))
            new_props = {}
            new_required = []
            for pname, pdef in old_props.items():
                pdef = dict(pdef)
                if pname in _GEMMA4_COLLIDING_PARAMS:
                    new_name = _GEMMA4_RENAME_PREFIX + pname
                else:
                    new_name = pname
                if not pdef.get("description"):
                    label = "REQUIRED. " if pname in required else ""
                    pdef["description"] = (
                        f"{label}The '{pname}' value"
                        f" (type: {pdef.get('type', 'string')})"
                    )
                new_props[new_name] = pdef
                new_required.append(new_name if pname in required else None)
            params["properties"] = new_props
            params["required"] = [r for r in new_required if r]
            func["parameters"] = params
        tool["function"] = func
        enriched.append(tool)
    return enriched


def restore_gemma4_param_names(arguments: dict) -> dict:
    """Reverse the parameter renaming done by :func:`enrich_tool_params_for_gemma4`."""
    restored = {}
    for k, v in arguments.items():
        if k.startswith(_GEMMA4_RENAME_PREFIX):
            original = k[len(_GEMMA4_RENAME_PREFIX):]
            if original in _GEMMA4_COLLIDING_PARAMS:
                restored[original] = v
                continue
        restored[k] = v
    return restored


def format_tool_call_for_message(tool_call: ToolCall) -> dict:
    """
    Format a ToolCall object for inclusion in a message.

    Args:
        tool_call: ToolCall object

    Returns:
        Dict representation suitable for message content
    """
    return {
        "id": tool_call.id,
        "type": tool_call.type,
        "function": {
            "name": tool_call.function.name,
            "arguments": tool_call.function.arguments,
        },
    }


# =============================================================================
# Structured Output (JSON Schema) Utilities
# =============================================================================


def validate_json_schema(
    data: Any, schema: Dict[str, Any]
) -> Tuple[bool, Optional[str]]:
    """
    Validate JSON data against a JSON Schema.

    Args:
        data: The JSON data to validate (dict, list, etc.)
        schema: JSON Schema specification

    Returns:
        Tuple of (is_valid, error_message)
        - is_valid: True if data matches schema
        - error_message: Error description if invalid, None if valid
    """
    try:
        validate(instance=data, schema=schema)
        return True, None
    except ValidationError as e:
        return False, str(e.message)


def extract_json_from_text(text: str) -> Optional[Dict[str, Any]]:
    """
    Extract JSON from model output text.

    Tries multiple strategies:
    1. Parse entire text as JSON
    2. Extract JSON from markdown code blocks
    3. Find JSON object/array in text

    Args:
        text: Raw model output text

    Returns:
        Parsed JSON data, or None if no valid JSON found
    """
    text = text.strip()

    # Strategy 1: Try to parse entire text as JSON
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Strategy 2: Extract from markdown code blocks
    # Match ```json ... ``` or ``` ... ```
    code_block_pattern = r"```(?:json)?\s*([\s\S]*?)\s*```"
    matches = re.findall(code_block_pattern, text)
    for match in matches:
        try:
            return json.loads(match.strip())
        except json.JSONDecodeError:
            continue

    # Strategy 3: Find JSON object or array in text
    # Look for { ... } or [ ... ]
    json_patterns = [
        r"(\{[\s\S]*\})",  # Object
        r"(\[[\s\S]*\])",  # Array
    ]
    for pattern in json_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                continue

    return None


def parse_json_output(
    text: str, response_format: Optional[Union[ResponseFormat, Dict[str, Any]]] = None
) -> Tuple[str, Optional[Dict[str, Any]], bool, Optional[str]]:
    """
    Parse JSON from model output when response_format is set.

    Args:
        text: Raw model output text
        response_format: ResponseFormat specification (optional)
            - If type="json_object", extracts any valid JSON
            - If type="json_schema", extracts and validates against schema

    Returns:
        Tuple of (cleaned_text, parsed_json, is_valid, error_message)
        - cleaned_text: Original text (preserved for reference)
        - parsed_json: Extracted JSON data, or None if extraction failed
        - is_valid: True if JSON is valid (and matches schema if specified)
        - error_message: Error description if invalid, None if valid
    """
    # Handle None or text format - just return original
    if response_format is None:
        return text, None, True, None

    # Normalize response_format to dict
    if isinstance(response_format, ResponseFormat):
        rf_dict = {"type": response_format.type, "json_schema": None}
        if response_format.json_schema:
            rf_dict["json_schema"] = {
                "name": response_format.json_schema.name,
                "description": response_format.json_schema.description,
                "schema": response_format.json_schema.schema_,
                "strict": response_format.json_schema.strict,
            }
    else:
        rf_dict = response_format

    format_type = rf_dict.get("type", "text")

    # text format - no JSON extraction
    if format_type == "text":
        return text, None, True, None

    # json_object or json_schema - extract JSON
    parsed = extract_json_from_text(text)

    if parsed is None:
        return text, None, False, "Failed to extract valid JSON from output"

    # json_object - just verify it's valid JSON (already done by extraction)
    if format_type == "json_object":
        return text, parsed, True, None

    # json_schema - validate against schema
    if format_type == "json_schema":
        json_schema_spec = rf_dict.get("json_schema") or {}
        schema = json_schema_spec.get("schema", {})

        if schema:
            is_valid, error = validate_json_schema(parsed, schema)
            if not is_valid:
                return text, parsed, False, f"JSON Schema validation failed: {error}"

        return text, parsed, True, None

    # Unknown format type - treat as text
    return text, None, True, None


def build_json_system_prompt(
    response_format: Optional[Union[ResponseFormat, Dict[str, Any]]] = None,
) -> Optional[str]:
    """
    Build a system prompt instruction for JSON output.

    For models without native JSON mode support, this adds instructions
    to the prompt to encourage proper JSON formatting.

    Args:
        response_format: ResponseFormat specification

    Returns:
        System prompt instruction string, or None if not needed
    """
    if response_format is None:
        return None

    # Normalize to dict
    if isinstance(response_format, ResponseFormat):
        rf_dict = {"type": response_format.type, "json_schema": None}
        if response_format.json_schema:
            rf_dict["json_schema"] = {
                "name": response_format.json_schema.name,
                "description": response_format.json_schema.description,
                "schema": response_format.json_schema.schema_,
                "strict": response_format.json_schema.strict,
            }
    else:
        rf_dict = response_format

    format_type = rf_dict.get("type", "text")

    if format_type == "text":
        return None

    if format_type == "json_object":
        return (
            "You must respond with valid JSON only. "
            "Do not include any explanation or text outside the JSON object."
        )

    if format_type == "json_schema":
        json_schema_spec = rf_dict.get("json_schema") or {}
        schema = json_schema_spec.get("schema", {})
        name = json_schema_spec.get("name", "response")
        description = json_schema_spec.get("description", "")

        prompt = f"You must respond with valid JSON matching the '{name}' schema."
        if description:
            prompt += f" {description}"
        prompt += (
            f"\n\nJSON Schema:\n```json\n{json.dumps(schema, indent=2)}\n```\n\n"
            "Respond with only the JSON object, no additional text or explanation."
        )
        return prompt

    return None


# ===== FILE: omlx/api/utils.py =====

# SPDX-License-Identifier: Apache-2.0
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
"""
Utility functions for text processing.
"""

import json
import re
from typing import Any, List

from .openai_models import Message

# Model families whose chat templates consume message.reasoning_content directly.
_NATIVE_REASONING_MODEL_TYPES = {
    "minimax_m3",
    "minimax_m3_vl",
    # Inkling's chat template renders history reasoning_content back into
    # <|content_thinking|> blocks.
    "inkling",
    "inkling_mm_model",
}


def uses_native_reasoning_content(
    model_name: str | None = None,
    *,
    config_model_type: str | None = None,
    engine_model_type: str | None = None,
    preserve_thinking_default: bool | None = None,
) -> bool:
    """Return whether history should keep reasoning in message fields."""
    if preserve_thinking_default is True:
        return True

    if config_model_type in _NATIVE_REASONING_MODEL_TYPES:
        return True
    if engine_model_type in _NATIVE_REASONING_MODEL_TYPES:
        return True

    lowered = (model_name or "").lower()
    return "minimax" in lowered and "m3" in lowered


# =============================================================================
# Partial Mode Detection
# =============================================================================


def detect_and_strip_partial(messages: list[dict]) -> bool:
    """Check if the final assistant message has partial=True; strip the field from all messages.

    Partial mode signals that the model should continue from the final assistant
    message rather than starting a new turn.  The ``partial`` key is not part of
    the chat-template contract, so it is always removed before the messages are
    passed to ``apply_chat_template``.

    Args:
        messages: List of message dicts (mutated in-place).

    Returns:
        True if the final message is an assistant message with ``partial=True``.
    """
    is_partial = (
        bool(messages)
        and messages[-1].get("role") == "assistant"
        and messages[-1].get("partial", False)
    )
    for msg in messages:
        msg.pop("partial", None)
    return is_partial


# =============================================================================
# Special Token Patterns
# =============================================================================

# Pattern to match special tokens that should be removed from output
SPECIAL_TOKENS_PATTERN = re.compile(
    r"<\|im_end\|>|<\|im_start\|>|<\|endoftext\|>|"
    r"<\|end\|>|<\|eot_id\|>|<\|start_header_id\|>|<\|end_header_id\|>|"
    r"<\|image\|>|<\|audio\|>|"  # Gemma 4 VLM special tokens
    r"\[e~\[|\]~b\]|\]~!b\[|\]!p~\[|\]!d~\[|"  # MiniMax M3 special tokens
    r"</s>|<s>|<pad>|\[PAD\]|\[SEP\]|\[CLS\]|"
    r"<eos>|<bos>|<end_of_turn>|<start_of_turn>"  # Gemma special tokens (fixes #1087)
)


def clean_special_tokens(text: str) -> str:
    """Clean model output by removing only special tokens.

    Preserves <think>...</think> blocks for downstream processing.

    Args:
        text: Raw model output

    Returns:
        Text with special tokens removed but think tags preserved
    """
    if not text:
        return text
    return SPECIAL_TOKENS_PATTERN.sub("", text).strip()


def remove_special_tokens_preserve_whitespace(text: str) -> str:
    """Remove special tokens without trimming surrounding whitespace."""
    if not text:
        return text
    return SPECIAL_TOKENS_PATTERN.sub("", text)


def clean_output_text(text: str) -> str:
    """Clean model output by removing special tokens and thinking blocks.

    Args:
        text: Raw model output

    Returns:
        Cleaned text with special tokens and <think> blocks removed
    """
    if not text:
        return text
    text = SPECIAL_TOKENS_PATTERN.sub("", text)
    from .thinking import extract_thinking

    _, content = extract_thinking(text)
    return content.strip()


# =============================================================================
# Text Content Extraction
# =============================================================================


def _extract_text_from_content_list(content: list) -> str:
    """Extract text parts from a content array, dropping non-text items.

    Handles content arrays from both OpenAI and Anthropic formats.
    Only items with type="text" are extracted; all others (tool_use,
    image, image_url, thinking, refusal, etc.) are silently dropped.
    """
    text_parts = []
    for item in content:
        # Convert Pydantic models to dict
        if hasattr(item, "model_dump"):
            item = item.model_dump()
        elif hasattr(item, "dict"):
            item = item.dict()

        if isinstance(item, dict):
            if item.get("type") == "text":
                text_parts.append(item.get("text", ""))
        elif isinstance(item, str):
            # Direct string in content list
            text_parts.append(item)

    return "\n".join(text_parts) if text_parts else ""


def _extract_multimodal_content_list(content: list) -> list:
    """Extract text, image, and audio parts from a content array.

    Keeps text, image_url, and input_audio items for VLM processing.
    Other content types (tool_use, thinking, refusal, etc.) are dropped.
    """
    parts = []
    for item in content:
        if hasattr(item, "model_dump"):
            item = item.model_dump()
        elif hasattr(item, "dict"):
            item = item.dict()
        if isinstance(item, dict):
            item_type = item.get("type")
            if item_type in ("text", "input_text"):
                text = item.get("text") or item.get("content") or ""
                parts.append({"type": "text", "text": text})
            elif item_type == "image_url":
                image_url_value = item.get("image_url")
                url = None
                if isinstance(image_url_value, str):
                    url = image_url_value
                elif isinstance(image_url_value, dict):
                    url = image_url_value.get("url")
                if url:
                    parts.append(
                        {
                            "type": "image_url",
                            "image_url": {"url": url},
                        }
                    )
            elif item_type == "input_image":
                image_url_value = item.get("image_url", item.get("input_image"))
                url = None
                if isinstance(image_url_value, str):
                    url = image_url_value
                elif isinstance(image_url_value, dict):
                    url = image_url_value.get("url")
                if url:
                    parts.append(
                        {
                            "type": "image_url",
                            "image_url": {"url": url},
                        }
                    )
            elif item_type == "image":
                # Anthropic format: convert to OpenAI image_url format
                source = item.get("source", {})
                if source.get("type") == "base64":
                    media_type = source.get("media_type", "image/jpeg")
                    data = source.get("data", "")
                    parts.append(
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{data}",
                            },
                        }
                    )
            elif item_type == "input_audio":
                # OpenAI audio format: pass through for engine-side decoding
                input_audio = item.get("input_audio")
                if input_audio and isinstance(input_audio, dict):
                    parts.append(
                        {
                            "type": "input_audio",
                            "input_audio": input_audio,
                        }
                    )
    return parts


# Roles eligible for merging when consecutive.
# System and tool messages are excluded: system messages have distinct semantics
# (e.g., JSON schema instructions), and tool messages carry tool_call_id.
_MERGEABLE_ROLES = {"user", "assistant"}
_PRESERVE_BOUNDARY_KEY = "_preserve_role_boundary"

# Match `role == "tool"` / `role == 'tool'` in a chat template.
_TOOL_ROLE_CHECK_RE = re.compile(r"==\s*['\"]tool['\"]")

_MID_SYSTEM_USER_MARKER = "__OMLX_MID_SYSTEM_PROBE_USER__"
_MID_SYSTEM_MARKER = "__OMLX_MID_SYSTEM_PROBE_SYSTEM__"
_MID_SYSTEM_ASSISTANT_MARKER = "__OMLX_MID_SYSTEM_PROBE_ASSISTANT__"
_MID_SYSTEM_PROBE_TOOL = [
    {
        "type": "function",
        "function": {
            "name": "omlx_probe_tool",
            "description": "oMLX chat-template probe tool.",
            "parameters": {"type": "object", "properties": {}},
        },
    }
]
_MID_SYSTEM_PROBE_CACHE: dict[tuple[Any, ...], bool] = {}


def _chat_template_supports_tool_role(tokenizer: Any) -> bool:
    """Check whether the tokenizer's chat template renders tool messages natively.

    mlx-lm / mlx-vlm only set ``has_tool_calling`` when their marker-based
    ``_infer_tool_parser`` recognises the chat template (qwen3_coder, json_tools,
    gemma4, etc.). Templates that branch on ``role == "tool"`` and render
    ``tool_calls`` but don't match any known marker (Qwen3 VL variants, custom
    fine-tunes) get flattened to ``role: "user"`` — making the model treat tool
    output as user instructions and breaking multi-turn tool flows (#1290).

    Strict superset of ``has_tool_calling``: if the tokenizer already flags
    itself, return True immediately. Otherwise probe the chat_template string
    for both a ``role == "tool"`` equality and the ``tool_calls`` variable —
    both together keep false positives down (a stray ``"tool"`` literal in a
    comment isn't enough).
    """
    if getattr(tokenizer, "has_tool_calling", False):
        return True
    chat_template = getattr(tokenizer, "chat_template", None)
    if not isinstance(chat_template, str):
        return False
    if not _TOOL_ROLE_CHECK_RE.search(chat_template):
        return False
    return "tool_calls" in chat_template


def _freeze_template_value(value: Any) -> Any:
    """Convert chat-template kwargs into a hashable cache-key value."""
    if isinstance(value, dict):
        return tuple(
            sorted((str(k), _freeze_template_value(v)) for k, v in value.items())
        )
    if isinstance(value, (list, tuple)):
        return tuple(_freeze_template_value(v) for v in value)
    if isinstance(value, set):
        return tuple(sorted(_freeze_template_value(v) for v in value))
    try:
        hash(value)
    except TypeError:
        return repr(value)
    return value


def _mid_system_probe_cache_key(
    tokenizer: Any,
    *,
    has_tools: bool,
    chat_template_kwargs: dict[str, Any] | None,
    placement: str,
    is_partial: bool,
) -> tuple[Any, ...]:
    chat_template = getattr(tokenizer, "chat_template", None)
    if isinstance(chat_template, str):
        template_fingerprint: Any = hash(chat_template)
    else:
        template_fingerprint = repr(chat_template)
    return (
        id(tokenizer),
        template_fingerprint,
        has_tools,
        _freeze_template_value(chat_template_kwargs or {}),
        placement,
        is_partial,
    )


def _apply_mid_system_probe_template(
    tokenizer: Any,
    probe_messages: list[dict],
    *,
    has_tools: bool,
    chat_template_kwargs: dict[str, Any] | None,
    is_partial: bool,
) -> str:
    template_kwargs: dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": not is_partial,
    }
    if is_partial:
        template_kwargs["continue_final_message"] = True
    if has_tools:
        template_kwargs["tools"] = _MID_SYSTEM_PROBE_TOOL
    if chat_template_kwargs:
        template_kwargs.update(chat_template_kwargs)

    try:
        rendered = tokenizer.apply_chat_template(probe_messages, **template_kwargs)
    except TypeError:
        if chat_template_kwargs:
            for key in chat_template_kwargs:
                template_kwargs.pop(key, None)
        template_kwargs.pop("tools", None)
        template_kwargs.pop("enable_thinking", None)
        rendered = tokenizer.apply_chat_template(probe_messages, **template_kwargs)

    if isinstance(rendered, str):
        return rendered
    if isinstance(rendered, list):
        return " ".join(str(token) for token in rendered)
    return str(rendered)


def chat_template_preserves_mid_system(
    tokenizer: Any | None,
    *,
    tools: list[dict] | None = None,
    chat_template_kwargs: dict[str, Any] | None = None,
    placement: str = "tail",
    is_partial: bool = False,
) -> bool:
    """Return whether the chat template renders a mid-system message in-place.

    This does not prove model-level semantics. It only verifies that the
    current tokenizer template keeps the system content after the preceding
    user turn instead of raising, dropping it, or moving it to the front.
    """
    if tokenizer is None or not hasattr(tokenizer, "apply_chat_template"):
        return False
    if placement not in {"tail", "between"}:
        return False

    has_tools = bool(tools)
    cache_key = _mid_system_probe_cache_key(
        tokenizer,
        has_tools=has_tools,
        chat_template_kwargs=chat_template_kwargs,
        placement=placement,
        is_partial=is_partial,
    )
    cached = _MID_SYSTEM_PROBE_CACHE.get(cache_key)
    if cached is not None:
        return cached

    probe_messages = [
        {"role": "user", "content": _MID_SYSTEM_USER_MARKER},
        {"role": "system", "content": _MID_SYSTEM_MARKER},
    ]
    if placement == "between":
        probe_messages.append(
            {"role": "assistant", "content": _MID_SYSTEM_ASSISTANT_MARKER}
        )

    try:
        rendered = _apply_mid_system_probe_template(
            tokenizer,
            probe_messages,
            has_tools=has_tools,
            chat_template_kwargs=chat_template_kwargs,
            is_partial=is_partial,
        )
    except Exception:
        _MID_SYSTEM_PROBE_CACHE[cache_key] = False
        return False

    user_idx = rendered.find(_MID_SYSTEM_USER_MARKER)
    system_idx = rendered.find(_MID_SYSTEM_MARKER)
    assistant_idx = rendered.find(_MID_SYSTEM_ASSISTANT_MARKER)

    supported = user_idx >= 0 and system_idx > user_idx
    if placement == "between":
        supported = supported and assistant_idx > system_idx

    _MID_SYSTEM_PROBE_CACHE[cache_key] = supported
    return supported


def _system_content_as_text(content: Any) -> str:
    if isinstance(content, list):
        return _extract_text_from_content_list(content)
    return content if isinstance(content, str) else str(content)


def _is_system_role(role: Any) -> bool:
    return role in {"system", "developer"}


def _merge_consecutive_system_messages(messages: list[dict]) -> list[dict]:
    """Merge adjacent system messages in-place without moving their position."""
    merged: list[dict] = []
    i = 0
    while i < len(messages):
        msg = messages[i]
        if not _is_system_role(msg.get("role")):
            merged.append(msg)
            i += 1
            continue

        parts: list[str] = []
        while i < len(messages) and _is_system_role(messages[i].get("role")):
            content = messages[i].get("content", "")
            if content:
                text = _system_content_as_text(content)
                if text:
                    parts.append(text)
            i += 1

        if parts:
            merged.append({"role": "system", "content": "\n\n".join(parts)})

    return merged


def _mid_system_placement_kinds(messages: list[dict]) -> set[str] | None:
    """Classify supported cache-preserving mid-system placements.

    Returns None when any non-leading system run has an unsupported position.
    """
    placements: set[str] = set()
    seen_non_system = False
    i = 0
    while i < len(messages):
        role = messages[i].get("role")
        if not _is_system_role(role):
            seen_non_system = True
            i += 1
            continue

        start = i
        while i < len(messages) and _is_system_role(messages[i].get("role")):
            i += 1

        if not seen_non_system:
            continue

        prev_role = messages[start - 1].get("role") if start > 0 else None
        next_role = messages[i].get("role") if i < len(messages) else None
        if prev_role != "user":
            return None
        if next_role is None:
            placements.add("tail")
        elif next_role == "assistant":
            placements.add("between")
        else:
            return None

    return placements


def has_nonleading_system_message(messages: list[dict]) -> bool:
    """Return True when a system message appears after a non-system turn."""
    seen_non_system = False
    for msg in messages:
        if _is_system_role(msg.get("role")):
            if seen_non_system:
                return True
        else:
            seen_non_system = True
    return False


def _is_text_only_content_list(content: Any) -> bool:
    if not isinstance(content, list):
        return False
    for part in content:
        if not isinstance(part, dict):
            return False
        if part.get("type", "text") != "text":
            return False
        text = part.get("text", "")
        if text is not None and not isinstance(text, str):
            return False
    return True


def _is_safe_user_note_target(msg: dict | None) -> bool:
    if not msg or msg.get("role") != "user":
        return False
    if msg.get(_PRESERVE_BOUNDARY_KEY):
        return False
    if msg.get("tool_calls") or msg.get("tool_call_id") or msg.get("tool_responses"):
        return False
    content = msg.get("content", "")
    return (
        content is None
        or isinstance(content, str)
        or _is_text_only_content_list(content)
    )


def _message_has_tool_calls(msg: dict | None) -> bool:
    return bool(msg and msg.get("role") == "assistant" and msg.get("tool_calls"))


def _format_system_note(parts: list[str]) -> str:
    return "[System note]\n" + "\n\n".join(parts) + "\n[/System note]"


def _merge_note_text(existing: str, note: str, *, placement: str) -> str:
    if not existing:
        return note
    if placement == "prepend":
        return f"{note}\n\n{existing}"
    return f"{existing}\n\n{note}"


def _rewrite_user_content_with_note(
    msg: dict,
    note: str,
    *,
    placement: str,
) -> dict:
    rewritten = dict(msg)
    content = rewritten.get("content", "")
    if isinstance(content, list):
        parts = [dict(part) for part in content]
        if not parts:
            rewritten["content"] = [{"type": "text", "text": note}]
            return rewritten
        index = 0 if placement == "prepend" else len(parts) - 1
        existing = parts[index].get("text") or ""
        parts[index]["text"] = _merge_note_text(
            existing,
            note,
            placement=placement,
        )
        rewritten["content"] = parts
        return rewritten

    existing = content if isinstance(content, str) else ""
    rewritten["content"] = _merge_note_text(existing, note, placement=placement)
    return rewritten


def _downgrade_mid_system_to_user_notes(messages: list[dict]) -> list[dict] | None:
    """Move unsupported non-leading system runs into adjacent safe user text.

    This keeps the native chat template/tool rendering path intact while making
    volatile tail notes cache-friendly. It deliberately refuses tool-call
    boundaries and multimodal user content, where changing roles is too risky.
    """
    rewritten: list[dict] = []
    seen_non_system = False
    i = 0

    while i < len(messages):
        msg = messages[i]
        if not _is_system_role(msg.get("role")):
            rewritten.append(msg)
            seen_non_system = True
            i += 1
            continue

        start = i
        parts: list[str] = []
        while i < len(messages) and _is_system_role(messages[i].get("role")):
            content = messages[i].get("content", "")
            if content:
                text = _system_content_as_text(content)
                if text:
                    parts.append(text)
            i += 1

        if not seen_non_system:
            rewritten.extend(messages[start:i])
            continue
        if not parts:
            continue

        note = _format_system_note(parts)
        next_msg = messages[i] if i < len(messages) else None
        next_role = next_msg.get("role") if next_msg is not None else None

        if _is_safe_user_note_target(rewritten[-1] if rewritten else None) and (
            next_msg is None or next_role == "assistant"
        ):
            rewritten[-1] = _rewrite_user_content_with_note(
                rewritten[-1],
                note,
                placement="append",
            )
            continue

        if _is_safe_user_note_target(next_msg):
            if _message_has_tool_calls(rewritten[-1] if rewritten else None):
                return None
            rewritten.append(
                _rewrite_user_content_with_note(
                    next_msg,
                    note,
                    placement="prepend",
                )
            )
            seen_non_system = True
            i += 1
            continue

        return None

    return rewritten


def prepare_system_messages_for_template(
    messages: list[dict],
    tokenizer: Any | None,
    *,
    tools: list[dict] | None = None,
    chat_template_kwargs: dict[str, Any] | None = None,
    is_partial: bool = False,
    merge_consecutive_roles: bool = True,
    unsupported_mid_system_policy: str = "strict",
) -> list[dict]:
    """Preserve cache-friendly mid-system turns when the template supports them.

    Unsupported placements or templates fall back to the historical behavior:
    all system messages are consolidated at the front.
    """
    messages = [dict(msg) for msg in messages]
    if unsupported_mid_system_policy not in {"strict", "user_note_safe"}:
        unsupported_mid_system_policy = "strict"

    def strict_fallback() -> list[dict]:
        prepared = _consolidate_system_messages(messages)
        if merge_consecutive_roles:
            prepared = _merge_consecutive_roles(prepared)
        return prepared

    def unsupported_fallback() -> list[dict]:
        if unsupported_mid_system_policy == "user_note_safe":
            prepared = _downgrade_mid_system_to_user_notes(messages)
            if prepared is not None:
                # _downgrade preserves leading system blocks as-is; merge
                # consecutive system messages so strict templates (Qwen3.6+)
                # that require a single leading system message don't fail.
                prepared = _merge_consecutive_system_messages(prepared)
                if merge_consecutive_roles:
                    prepared = _merge_consecutive_roles(prepared)
                return prepared
        return strict_fallback()

    placements = _mid_system_placement_kinds(messages)
    if not placements:
        if placements is None:
            return unsupported_fallback()
        return _merge_consecutive_system_messages(messages)

    if is_partial:
        return strict_fallback()

    can_preserve = all(
        chat_template_preserves_mid_system(
            tokenizer,
            tools=tools,
            chat_template_kwargs=chat_template_kwargs,
            placement=placement,
            is_partial=is_partial,
        )
        for placement in placements
    )
    if can_preserve:
        return _merge_consecutive_system_messages(messages)

    return unsupported_fallback()


def _drop_void_assistant_messages(messages: list[dict]) -> list[dict]:
    """Drop assistant messages that have no content and no tool_calls.

    Strict chat templates (e.g., Devstral/Mistral) raise an error when an
    assistant message has empty content and no tool_calls.  These void messages
    carry no information and can appear when a client echoes back a response
    that had only tool calls which were not preserved in its history.

    Messages with ``tool_responses`` (Gemma 4 format) or ``reasoning_content``
    (Qwen 3.6+ native reasoning field) are never dropped even when content is
    empty — they carry their own payload the template renders.
    """
    return [
        msg
        for msg in messages
        if not (
            msg.get("role") == "assistant"
            and not msg.get("content")
            and not msg.get("tool_calls")
            and not msg.get("tool_responses")
            and not msg.get("reasoning_content")
        )
    ]


def _consolidate_system_messages(messages: list[dict]) -> list[dict]:
    """Move all system messages to the front, merged into one.

    Models with strict chat templates (e.g., Qwen3.5) require the system
    message to appear first.  Clients may send system or developer messages
    mid-conversation, so we consolidate them defensively.
    """
    system_parts: list[str] = []
    non_system: list[dict] = []
    for msg in messages:
        if _is_system_role(msg.get("role")):
            content = msg.get("content", "")
            if content:
                if isinstance(content, list):
                    text = _extract_text_from_content_list(content)
                    if text:
                        system_parts.append(text)
                else:
                    system_parts.append(content)
        else:
            non_system.append(msg)

    if not system_parts:
        return messages

    merged_system = {"role": "system", "content": "\n\n".join(system_parts)}
    return [merged_system] + non_system


def _merge_consecutive_roles(messages: list[dict]) -> list[dict]:
    """Merge consecutive messages with the same mergeable role.

    Models with strict chat templates (e.g., Gemma-3) enforce alternating
    user/assistant roles and reject consecutive same-role messages.
    OpenAI's API accepts these, so we merge them for compatibility.

    Args:
        messages: List of processed message dicts with 'role' and 'content'.

    Returns:
        New list with consecutive same-role messages merged using "\\n\\n".
    """
    if not messages:
        return messages

    merged: list[dict] = [messages[0].copy()]

    for msg in messages[1:]:
        prev = merged[-1]
        if (
            msg["role"] == prev["role"]
            and msg["role"] in _MERGEABLE_ROLES
            and not prev.get(_PRESERVE_BOUNDARY_KEY)
            and not msg.get(_PRESERVE_BOUNDARY_KEY)
        ):
            prev_content = prev.get("content", "")
            new_content = msg.get("content", "")
            if prev_content and new_content:
                prev_is_list = isinstance(prev_content, list)
                new_is_list = isinstance(new_content, list)
                if prev_is_list or new_is_list:
                    # Convert both to list form for safe concatenation
                    if not prev_is_list:
                        prev_content = [{"type": "text", "text": prev_content}]
                    if not new_is_list:
                        new_content = [{"type": "text", "text": new_content}]
                    prev["content"] = prev_content + new_content
                else:
                    prev["content"] = prev_content + "\n\n" + new_content
            elif new_content:
                prev["content"] = new_content
        else:
            merged.append(msg.copy())

    return merged


def _apply_reasoning_reconstruction(
    role: str,
    content: Any,
    reasoning: str | None,
    native: bool,
) -> tuple[Any, str | None]:
    """Reconstruct reasoning on a historical assistant message.

    External clients echo reasoning back via the OpenAI ``reasoning_content``
    field (or Anthropic ``thinking`` blocks).  Chat templates fall into two
    camps:

    * ``native=True`` — template understands ``message.reasoning_content``
      as a top-level field (Qwen 3.6+).  Content stays clean and reasoning
      travels separately.
    * ``native=False`` — template only parses ``<think>...</think>`` embedded
      in content.  Reasoning is inlined into content as a fallback.

    Returns ``(new_content, reasoning_out)`` where ``reasoning_out`` is the
    string to attach as a ``reasoning_content`` field, or ``None`` to skip.
    """
    if role != "assistant" or not reasoning:
        if role != "assistant" or not native:
            return content, None
        text = content if isinstance(content, str) else ""
        if isinstance(content, list):
            text = _extract_text_from_content_list(content)
        from .thinking import extract_thinking

        inline_reasoning, inline_content = extract_thinking(text)
        if inline_reasoning:
            return inline_content, inline_reasoning
        return content, None
    text = content if isinstance(content, str) else ""
    if isinstance(content, list):
        text = _extract_text_from_content_list(content)
    if native:
        return text, reasoning
    return f"<think>\n{reasoning}\n</think>\n\n{text}", None


def extract_text_content(
    messages: List[Message],
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    native_reasoning_content: bool = False,
    consolidate_system_messages: bool = True,
) -> List[dict]:
    """
    Extract text content from OpenAI-format messages.

    Handles:
    - Simple text messages
    - Content arrays (extracts text parts only)
    - Tool call messages (assistant with tool_calls)
    - Tool response messages (role="tool")

    Args:
        messages: List of Message objects
        max_tool_result_tokens: Maximum token count for tool results.
        tokenizer: Tokenizer instance for token counting and truncation.
        native_reasoning_content: If True, pass ``reasoning_content`` through
            as a message-level field (Qwen 3.6+ templates).  If False, inline
            ``<think>...</think>`` into content as a fallback.
        consolidate_system_messages: If True, preserve historical strict-template
            behavior by moving system messages to the front. Server code can
            set this to False and call ``prepare_system_messages_for_template``
            after tools/template kwargs are known.

    Returns:
        List of {"role": str, "content": str}
    """
    processed_messages = []

    for msg in messages:
        role = msg.role
        content = msg.content

        # Reconstruct reasoning for historical assistant messages.  Native
        # mode passes reasoning as a separate field; fallback inlines it as
        # <think>...</think> in content.
        reasoning = getattr(msg, "reasoning_content", None)
        content, reasoning_out = _apply_reasoning_reconstruction(
            role, content, reasoning, native_reasoning_content
        )

        # Normalize "developer" role to "system" (OpenAI API compatibility)
        if role == "developer":
            role = "system"

        # Handle tool response messages (role="tool")
        if role == "tool":
            tool_call_id = getattr(msg, "tool_call_id", None) or ""
            # Convert list content to string if needed
            if isinstance(content, list):
                tool_content = _extract_text_from_content_list(content)
            else:
                tool_content = content if content else ""
            # Apply truncation if configured
            if max_tool_result_tokens and tokenizer and tool_content:
                from .anthropic_utils import truncate_tool_result

                tool_content = truncate_tool_result(
                    tool_content, max_tool_result_tokens, tokenizer
                )
            # Preserve structured format for models with native tool calling
            # so the chat template renders tool results in the model's native format
            if _chat_template_supports_tool_role(tokenizer):
                processed_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call_id,
                        "content": tool_content,
                    }
                )
            else:
                processed_messages.append(
                    {
                        "role": "user",  # mlx-lm expects user/assistant roles
                        "content": f"[Tool Result ({tool_call_id})]: {tool_content}",
                        _PRESERVE_BOUNDARY_KEY: True,
                    }
                )
            continue

        # Handle assistant messages with tool_calls
        if role == "assistant" and hasattr(msg, "tool_calls") and msg.tool_calls:
            if isinstance(content, list):
                content = _extract_text_from_content_list(content)
            msg_dict = {"role": role, "content": content if content else ""}
            if reasoning_out is not None:
                msg_dict["reasoning_content"] = reasoning_out
            if getattr(msg, "name", None):
                msg_dict["name"] = msg.name
            if getattr(msg, "partial", False):
                msg_dict["partial"] = True

            # Preserve structured tool_calls for models with native tool calling
            # so the chat template renders them in the model's native format.
            # Without this, models mimic text-formatted tool calls from history
            # instead of generating their native parseable format.
            if _chat_template_supports_tool_role(tokenizer):
                tool_calls_list = []
                for tc in msg.tool_calls:
                    if isinstance(tc, dict):
                        func = tc.get("function", {})
                        tool_calls_list.append(
                            {
                                "id": tc.get("id", ""),
                                "function": {
                                    "name": func.get("name", ""),
                                    "arguments": _try_parse_json(
                                        func.get("arguments", "{}")
                                    ),
                                },
                            }
                        )
                    else:
                        args_str = (
                            getattr(tc.function, "arguments", "{}")
                            if hasattr(tc, "function")
                            else "{}"
                        )
                        tool_calls_list.append(
                            {
                                "id": getattr(tc, "id", ""),
                                "function": {
                                    "name": (
                                        getattr(tc.function, "name", "")
                                        if hasattr(tc, "function")
                                        else ""
                                    ),
                                    "arguments": _try_parse_json(args_str),
                                },
                            }
                        )
                msg_dict["tool_calls"] = tool_calls_list
            else:
                # Text fallback for models without native tool calling
                tool_calls_text = []
                for tc in msg.tool_calls:
                    if isinstance(tc, dict):
                        func = tc.get("function", {})
                        name = func.get("name", "unknown")
                        args = func.get("arguments", "{}")
                        tool_calls_text.append(f"[Calling tool: {name}({args})]")
                text = msg_dict["content"]
                if tool_calls_text:
                    text = (text + "\n" if text else "") + "\n".join(tool_calls_text)
                msg_dict["content"] = text
            msg_dict[_PRESERVE_BOUNDARY_KEY] = True

            processed_messages.append(msg_dict)
            continue

        # Build optional extra fields from the source message
        _extra: dict = {}
        if getattr(msg, "name", None):
            _extra["name"] = msg.name
        if getattr(msg, "partial", False):
            _extra["partial"] = True
        if reasoning_out is not None:
            _extra["reasoning_content"] = reasoning_out

        # Handle None content
        if content is None:
            processed_messages.append({"role": role, "content": "", **_extra})
            continue

        if isinstance(content, str):
            # Simple text message
            processed_messages.append({"role": role, "content": content, **_extra})
        elif isinstance(content, list):
            # Content array - extract text parts only
            combined_text = _extract_text_from_content_list(content)
            processed_messages.append(
                {"role": role, "content": combined_text, **_extra}
            )
        else:
            # Unknown format, try to convert
            processed_messages.append({"role": role, "content": str(content), **_extra})

    processed_messages = _drop_void_assistant_messages(processed_messages)
    if consolidate_system_messages:
        processed_messages = _consolidate_system_messages(processed_messages)
    return _merge_consecutive_roles(processed_messages)


def extract_multimodal_content(
    messages: List[Message],
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    native_reasoning_content: bool = False,
    consolidate_system_messages: bool = True,
) -> List[dict]:
    """
    Extract content from messages, preserving image_url parts for VLM.

    Same as extract_text_content but keeps image_url content parts
    in their original list format for VLM processing.

    Args:
        messages: List of Message objects
        max_tool_result_tokens: Maximum token count for tool results.
        tokenizer: Tokenizer instance for token counting and truncation.
        native_reasoning_content: If True, pass ``reasoning_content`` through
            as a message-level field.  See ``extract_text_content``.
        consolidate_system_messages: See ``extract_text_content``.

    Returns:
        List of message dicts. Messages with images have content as list.
    """
    processed_messages = []

    for msg in messages:
        role = msg.role
        content = msg.content

        # Reconstruct reasoning (see extract_text_content).
        reasoning = getattr(msg, "reasoning_content", None)
        content, reasoning_out = _apply_reasoning_reconstruction(
            role, content, reasoning, native_reasoning_content
        )

        if role == "developer":
            role = "system"

        # Tool response messages - same as extract_text_content
        if role == "tool":
            tool_call_id = getattr(msg, "tool_call_id", None) or ""
            # Convert list content to string if needed
            if isinstance(content, list):
                tool_content = _extract_text_from_content_list(content)
            else:
                tool_content = content if content else ""
            if max_tool_result_tokens and tokenizer and tool_content:
                from .anthropic_utils import truncate_tool_result

                tool_content = truncate_tool_result(
                    tool_content, max_tool_result_tokens, tokenizer
                )
            if _chat_template_supports_tool_role(tokenizer):
                processed_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call_id,
                        "content": tool_content,
                    }
                )
            else:
                processed_messages.append(
                    {
                        "role": "user",
                        "content": f"[Tool Result ({tool_call_id})]: {tool_content}",
                        _PRESERVE_BOUNDARY_KEY: True,
                    }
                )
            continue

        # Assistant with tool_calls - same as extract_text_content
        if role == "assistant" and hasattr(msg, "tool_calls") and msg.tool_calls:
            if isinstance(content, list):
                content = _extract_text_from_content_list(content)
            msg_dict = {"role": role, "content": content if content else ""}
            if reasoning_out is not None:
                msg_dict["reasoning_content"] = reasoning_out
            if getattr(msg, "name", None):
                msg_dict["name"] = msg.name
            if getattr(msg, "partial", False):
                msg_dict["partial"] = True

            if _chat_template_supports_tool_role(tokenizer):
                tool_calls_list = []
                for tc in msg.tool_calls:
                    if isinstance(tc, dict):
                        func = tc.get("function", {})
                        tool_calls_list.append(
                            {
                                "id": tc.get("id", ""),
                                "function": {
                                    "name": func.get("name", ""),
                                    "arguments": _try_parse_json(
                                        func.get("arguments", "{}")
                                    ),
                                },
                            }
                        )
                    else:
                        args_str = (
                            getattr(tc.function, "arguments", "{}")
                            if hasattr(tc, "function")
                            else "{}"
                        )
                        tool_calls_list.append(
                            {
                                "id": getattr(tc, "id", ""),
                                "function": {
                                    "name": (
                                        getattr(tc.function, "name", "")
                                        if hasattr(tc, "function")
                                        else ""
                                    ),
                                    "arguments": _try_parse_json(args_str),
                                },
                            }
                        )
                msg_dict["tool_calls"] = tool_calls_list
            else:
                tool_calls_text = []
                for tc in msg.tool_calls:
                    if isinstance(tc, dict):
                        func = tc.get("function", {})
                        name = func.get("name", "unknown")
                        args = func.get("arguments", "{}")
                        tool_calls_text.append(f"[Calling tool: {name}({args})]")
                text = msg_dict["content"]
                if tool_calls_text:
                    text = (text + "\n" if text else "") + "\n".join(tool_calls_text)
                msg_dict["content"] = text
            msg_dict[_PRESERVE_BOUNDARY_KEY] = True

            processed_messages.append(msg_dict)
            continue

        # Build optional extra fields from the source message
        _extra: dict = {}
        if getattr(msg, "name", None):
            _extra["name"] = msg.name
        if getattr(msg, "partial", False):
            _extra["partial"] = True
        if reasoning_out is not None:
            _extra["reasoning_content"] = reasoning_out

        if content is None:
            processed_messages.append({"role": role, "content": "", **_extra})
            continue

        if isinstance(content, str):
            processed_messages.append({"role": role, "content": content, **_extra})
        elif isinstance(content, list):
            # Preserve image_url and input_audio parts for VLM processing
            multimodal_parts = _extract_multimodal_content_list(content)
            multimodal_types = {"image_url", "input_audio"}
            has_multimodal = any(
                p.get("type") in multimodal_types for p in multimodal_parts
            )
            if has_multimodal:
                # Keep as content list for VLM engine
                processed_messages.append(
                    {"role": role, "content": multimodal_parts, **_extra}
                )
            else:
                # Text-only, flatten to string
                combined_text = _extract_text_from_content_list(content)
                processed_messages.append(
                    {"role": role, "content": combined_text, **_extra}
                )
        else:
            processed_messages.append({"role": role, "content": str(content), **_extra})

    processed_messages = _drop_void_assistant_messages(processed_messages)
    if consolidate_system_messages:
        processed_messages = _consolidate_system_messages(processed_messages)
    return processed_messages


# =============================================================================
# Harmony (gpt-oss) Message Extraction
# =============================================================================


def _try_parse_json(s: str):
    """
    Try to parse a string as JSON. Returns parsed dict/list if valid JSON,
    otherwise returns the original string.

    This is needed because Harmony chat_template uses |tojson filter,
    which would double-encode strings that are already JSON.
    """
    if not isinstance(s, str):
        return s
    s = s.strip()
    if not s:
        return s
    # Quick check: must start with { or [ to be JSON object/array
    if not (s.startswith("{") or s.startswith("[")):
        return s
    try:
        return json.loads(s)
    except (json.JSONDecodeError, ValueError):
        return s


def _wrap_truncated_for_harmony(truncated_text: str) -> dict:
    """Wrap truncated tool result in a dict for Harmony |tojson compatibility.

    The Harmony chat_template applies |tojson to tool result content.
    When truncation breaks valid JSON, the content becomes a string, and
    |tojson would double-encode it (wrapping in quotes and escaping).
    This function wraps the truncated text in a dict so |tojson produces
    a clean JSON object instead.

    Args:
        truncated_text: Text with truncation notice appended.

    Returns:
        Dict with 'output' key containing the truncated content and
        'truncated' key with a human-readable summary.
    """
    match = re.search(
        r'\n\n<truncated total_tokens="(\d+)" shown_tokens="(\d+)" />\s*$',
        truncated_text,
    )
    if match:
        return {
            "output": truncated_text[: match.start()],
            "truncated": f"Showing {match.group(2)} of {match.group(1)} tokens",
        }
    return {"output": truncated_text}


def extract_harmony_messages(
    messages: list,
    max_tool_result_tokens: int | None = None,
    tokenizer: Any | None = None,
    consolidate_system_messages: bool = True,
) -> List[dict]:
    """
    Extract messages for Harmony (gpt-oss) models.

    Unlike extract_text_content(), this function preserves:
    - tool messages: role="tool" with tool_call_id (chat_template handles conversion)
    - assistant tool_calls: tool_calls field intact (chat_template handles conversion)

    The Harmony chat_template expects standard OpenAI format and converts:
    - role="tool" → <|start|>functions.{name} to=assistant<|channel|>commentary...
    - assistant.tool_calls → <|start|>assistant to=functions.{name}<|channel|>commentary...

    IMPORTANT: The chat_template uses |tojson filter on:
    - tool_call.arguments (line 299)
    - message.content for tool results (line 322)

    If these are already JSON strings, |tojson would double-encode them.
    So we parse JSON strings to dicts before passing to the template.

    Args:
        messages: List of Message objects
        max_tool_result_tokens: Maximum token count for tool results.
        tokenizer: Tokenizer instance for token counting and truncation.
        consolidate_system_messages: See ``extract_text_content``.

    Returns:
        List of message dicts with tool-related fields preserved
    """
    processed_messages = []

    # Normalize to plain dicts -- callers may pass Pydantic Message
    # objects (OpenAI path) or plain dicts (Anthropic path).
    raw: list[dict] = []
    for msg in messages:
        if hasattr(msg, "model_dump"):
            raw.append(msg.model_dump())
        elif isinstance(msg, dict):
            raw.append(dict(msg))
        else:
            d: dict = {
                "role": getattr(msg, "role", "user"),
                "content": getattr(msg, "content", ""),
            }
            tool_call_id = getattr(msg, "tool_call_id", None)
            if tool_call_id is not None:
                d["tool_call_id"] = tool_call_id
            tool_calls = getattr(msg, "tool_calls", None)
            if tool_calls is not None:
                d["tool_calls"] = tool_calls
            raw.append(d)

    for msg in raw:
        role = msg.get("role", "user")
        content = msg.get("content")

        # Normalize "developer" role to "system" (OpenAI API compatibility)
        if role == "developer":
            role = "system"

        # Tool response messages - preserve role and tool_call_id
        # Parse content as JSON if possible (chat_template applies |tojson)
        if role == "tool":
            # Convert list content to string if needed
            if isinstance(content, list):
                tool_content = _extract_text_from_content_list(content)
            else:
                tool_content = content if content else ""
            if max_tool_result_tokens and tokenizer and tool_content:
                from .anthropic_utils import truncate_tool_result

                # Parse JSON BEFORE truncation for better line-boundary cuts.
                # Harmony chat_template applies |tojson to content, so content
                # must be a dict (not a string) to avoid double-encoding.
                parsed_json = _try_parse_json(tool_content)
                if isinstance(parsed_json, (dict, list)):
                    # Valid JSON - pretty-print for line-boundary truncation
                    pretty = json.dumps(parsed_json, indent=2, ensure_ascii=False)
                    truncated = truncate_tool_result(
                        pretty, max_tool_result_tokens, tokenizer
                    )
                    if "<truncated " in truncated:
                        # Truncation broke JSON - wrap in dict for |tojson
                        parsed_content = _wrap_truncated_for_harmony(truncated)
                    else:
                        # Not truncated - use parsed dict/list
                        parsed_content = parsed_json
                else:
                    # Not JSON - truncate raw text, keep as string
                    parsed_content = truncate_tool_result(
                        tool_content, max_tool_result_tokens, tokenizer
                    )
            else:
                # No truncation configured - just parse JSON if possible
                parsed_content = _try_parse_json(tool_content)
            processed_messages.append(
                {
                    "role": "tool",
                    "tool_call_id": msg.get("tool_call_id", "") or "",
                    "content": parsed_content,
                }
            )
            continue

        # Assistant messages - preserve tool_calls field
        if role == "assistant":
            msg_dict = {"role": role}

            # Handle content (may be string or list)
            if content is None:
                msg_dict["content"] = ""
            elif isinstance(content, str):
                msg_dict["content"] = content
            elif isinstance(content, list):
                # Extract text parts from content array
                msg_dict["content"] = _extract_text_from_content_list(content)
            else:
                msg_dict["content"] = str(content)

            # Preserve tool_calls field for chat_template
            # Parse arguments as JSON if possible (chat_template applies |tojson)
            if msg.get("tool_calls"):
                tool_calls_list = []
                for tc in msg["tool_calls"]:
                    if isinstance(tc, dict):
                        args_str = tc.get("function", {}).get("arguments", "{}")
                        tool_calls_list.append(
                            {
                                "id": tc.get("id", ""),
                                "function": {
                                    "name": tc.get("function", {}).get("name", ""),
                                    "arguments": _try_parse_json(args_str),
                                },
                            }
                        )
                    else:
                        # Pydantic model
                        args_str = (
                            getattr(tc.function, "arguments", "{}")
                            if hasattr(tc, "function")
                            else "{}"
                        )
                        tool_calls_list.append(
                            {
                                "id": getattr(tc, "id", ""),
                                "function": {
                                    "name": (
                                        getattr(tc.function, "name", "")
                                        if hasattr(tc, "function")
                                        else ""
                                    ),
                                    "arguments": _try_parse_json(args_str),
                                },
                            }
                        )
                msg_dict["tool_calls"] = tool_calls_list
                msg_dict[_PRESERVE_BOUNDARY_KEY] = True

            processed_messages.append(msg_dict)
            continue

        # Other messages (user, system, developer)
        if content is None:
            processed_messages.append({"role": role, "content": ""})
        elif isinstance(content, str):
            processed_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            # Extract text parts from content array
            processed_messages.append(
                {"role": role, "content": _extract_text_from_content_list(content)}
            )
        else:
            processed_messages.append({"role": role, "content": str(content)})

    processed_messages = _drop_void_assistant_messages(processed_messages)
    if consolidate_system_messages:
        processed_messages = _consolidate_system_messages(processed_messages)
    return _merge_consecutive_roles(processed_messages)


# ===== FILE: omlx/cache/__init__.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Cache module - unified cache management for oMLX.

This package contains cache implementations including:
- Prefix caching for KV state reuse
- Paged cache for memory-efficient KV storage
- VLM cache for vision-language model optimization
- SSD cache for disk-based persistence
"""

# Stats
from .stats import (
    BaseCacheStats,
    PagedCacheStats,
    VLMCacheStats,
    PagedSSDCacheStats,
)

# Interfaces
from .interface import CacheManager

# Paged cache implementations
from .paged_cache import (
    PagedCacheManager,
    CacheBlock,
    BlockTable,
    FreeKVCacheBlockQueue,
    BlockHashToBlockMap,
    BlockHash,
    compute_block_hash,
)

# Prefix cache implementations (SSD-only)
from .prefix_cache import (
    BlockAwarePrefixCache,
    BlockCacheEntry,
)

# Paged SSD cache implementations
from .paged_ssd_cache import (
    PagedSSDCacheManager,
    PagedSSDBlockMetadata,
    PagedSSDCacheIndex,
    parse_size,
)

# Vision feature cache
from .vision_feature_cache import (
    VisionFeatureSSDCache,
    VisionFeatureSSDEntry,
)

# Managers
from .recovery import CacheRecoveryManager

# Factory
from .factory import CacheConfig, CacheFactory

# Type handlers
from .type_handlers import (
    CacheType,
    CacheTypeHandler,
    CacheStateInfo,
    KVCacheHandler,
    RotatingKVCacheHandler,
    ArraysCacheHandler,
    CacheListHandler,
    DefaultCacheHandler,
    SizedArraysCache,
)

# Type registry
from .type_registry import CacheTypeRegistry

# Hybrid cache config
from .hybrid_cache import (
    LayerCacheConfig,
    ModelCacheConfig,
    create_default_kvcache_config,
)

__all__ = [
    # Stats
    "BaseCacheStats",
    "PagedCacheStats",
    "VLMCacheStats",
    "PagedSSDCacheStats",
    # Interfaces
    "CacheManager",
    # Paged cache
    "PagedCacheManager",
    "CacheBlock",
    "BlockTable",
    "FreeKVCacheBlockQueue",
    "BlockHashToBlockMap",
    "BlockHash",
    "compute_block_hash",
    # Prefix cache (SSD-only)
    "BlockAwarePrefixCache",
    "BlockCacheEntry",
    # Paged SSD cache
    "PagedSSDCacheManager",
    "PagedSSDBlockMetadata",
    "PagedSSDCacheIndex",
    "parse_size",
    # Vision feature cache
    "VisionFeatureSSDCache",
    "VisionFeatureSSDEntry",
    # Managers
    "CacheRecoveryManager",
    # Factory
    "CacheConfig",
    "CacheFactory",
    # Type handlers
    "CacheType",
    "CacheTypeHandler",
    "CacheStateInfo",
    "KVCacheHandler",
    "RotatingKVCacheHandler",
    "ArraysCacheHandler",
    "CacheListHandler",
    "DefaultCacheHandler",
    "SizedArraysCache",
    # Type registry
    "CacheTypeRegistry",
    # Hybrid cache config
    "LayerCacheConfig",
    "ModelCacheConfig",
    "create_default_kvcache_config",
]


# ===== FILE: omlx/cache/_rotating_subclass.py =====

# SPDX-License-Identifier: Apache-2.0
"""RotatingKVCache subclass tailored for omlx's restored-cache contract.

mlx-lm v0.31.3's BatchRotatingKVCache.merge() does

    keys[i:i+1, :, p:p+l] = c._temporal_order(c.keys)[..., -l:, :]

where ``l = c.size()`` and the default ``RotatingKVCache.size()`` returns
``min(offset, max_size)`` ignoring the actual buffer length. For
SSD-restored caches with ``keys.shape[2] < max_size`` and
``offset >= max_size`` this overshoots and either (a) shape-mismatches
the RHS slice when ``keys.shape[2] < l``, or (b) (when omlx zero-pads
the buffer up to max_size) exposes zero positions to attention causing
softmax dilution that surfaces as infinite loops or empty content
(issues #934, #903, #900).

This subclass clamps ``size()`` to the actual buffer length so merge is
always well-defined without any zero-padding trick on omlx's side.
"""
from __future__ import annotations

from mlx_lm.models.cache import RotatingKVCache


class PrefillReadyRotatingKVCache(RotatingKVCache):
    """RotatingKVCache that reports actual buffer length from ``size()``.

    The default ``size()`` returns ``min(offset, max_size)`` which is the
    logical token count. For caches restored from SSD whose buffer was
    sliced shorter than ``max_size`` (e.g. extract() stripped left
    padding), the logical count can exceed ``keys.shape[2]``. mlx-lm's
    merge then either over-reads the RHS or, when omlx pre-pads with
    zeros, lets those zeros leak into attention.

    Clamping to ``keys.shape[2]`` keeps merge consistent: the row gets
    exactly ``keys.shape[2]`` real entries, padded on the left by the
    enclosing batch (via ``left_padding``) instead of by phantom zeros.
    """

    def size(self):
        if self.keys is None:
            return 0
        buffer_len = self.keys.shape[2]
        if buffer_len == 0:
            return 0
        return min(super().size(), buffer_len)


# ===== FILE: omlx/cache/boundary_snapshot_store.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Boundary Snapshot SSD Store for oMLX.

Stores non-sliceable cache layer snapshots (e.g. ArraysCache) to SSD during
prefill, freeing GPU memory immediately.  At request completion the snapshots
are loaded back one block at a time for final SSD cache storage.

Uses the same async-write pattern as PagedSSDCacheManager: tensors are
serialized to raw bytes on the inference thread (Metal-safe), buffered in
``_pending_writes`` for instant read-back, and flushed to disk by a
background writer thread via ``_write_safetensors_no_mx``.
"""

from __future__ import annotations

import json
import logging
import os
import queue
import shutil
import threading
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Any

from .paged_ssd_cache import (
    HAS_MLX,
    _encode_shape,
    _extract_tensor_bytes,
    _has_zero_dim,
    _restore_tensor_from_bytes,
    _write_safetensors_no_mx,
)
from .pooling_delta import compact_pooling_cache_snapshot

if HAS_MLX:
    import mlx.core as mx

logger = logging.getLogger(__name__)

# Max pending writes before save() blocks.
_MAX_PENDING_WRITES = 128


def reset_boundary_snapshot_root(base_dir: Path) -> None:
    """Remove all boundary snapshot sessions for a server lifecycle boundary."""
    snapshot_root = base_dir / "_boundary_snapshots"
    if snapshot_root.exists():
        try:
            shutil.rmtree(snapshot_root)
        except Exception as e:
            logger.warning("Failed to reset boundary snapshots: %s", e)
    snapshot_root.mkdir(parents=True, exist_ok=True)


class BoundarySnapshotSSDStore:
    """Temporary SSD storage for boundary cache snapshots.

    Stores ArraysCache/RotatingKVCache boundary snapshots to SSD during
    prefill to avoid GPU memory accumulation.  Files are ephemeral and
    cleaned up when the request completes or aborts.

    Parameters
    ----------
    base_dir : Path
        Parent directory for the SSD cache (typically ``paged_ssd_cache_dir``).
        Snapshots are stored under
        ``base_dir/_boundary_snapshots/<session_id>/<request_id>/``.
    """

    # Timeouts applied when acquiring _writer_busy from each cleanup
    # path. cleanup_request is called from the scheduler's abort hot
    # path (~3 sites) and must yield faster than cleanup_all, which
    # also runs at startup / reset where blocking longer is tolerable
    # in exchange for a stronger orphan-avoidance guarantee. The
    # worst-case impact on the timeout fallback is identical in both
    # paths — an orphan file in the recreated dir until the next
    # constructor cleanup — so the only knob is per-call latency.
    _CLEANUP_ALL_TIMEOUT_S = 5.0
    _CLEANUP_REQUEST_TIMEOUT_S = 2.0

    def __init__(self, base_dir: Path) -> None:
        self._snapshot_root = base_dir / "_boundary_snapshots"
        self._session_id = f"{os.getpid()}-{uuid.uuid4().hex}"
        self._snapshot_dir = self._snapshot_root / self._session_id
        self._snapshot_dir.mkdir(parents=True, exist_ok=True)

        # request_id -> {token_count -> file_path}
        self._file_registry: dict[str, dict[int, Path]] = {}
        self._registry_lock = threading.Lock()

        # Pending writes buffer — raw bytes for instant read-back.
        # key: (request_id, token_count)
        self._pending_writes: dict[tuple[str, int], dict] = {}
        self._pending_lock = threading.Lock()

        # Cancelled requests with remaining queue item counts. Writer
        # thread decrements on each skip; entry is deleted when count
        # reaches zero, preventing unbounded growth. All access is
        # guarded by ``_cancelled_lock`` — the dict was previously
        # mutated unlocked from cleanup_request, cleanup_all, and the
        # writer thread, creating lost-cancellation and counter-
        # underflow races.
        self._cancelled_requests: dict[str, int] = {}
        self._cancelled_lock = threading.Lock()

        # Background writer thread.
        self._write_queue: queue.Queue = queue.Queue(maxsize=_MAX_PENDING_WRITES)
        self._shutdown = threading.Event()
        # Held by the writer for the duration of each item's processing.
        # cleanup_all() acquires it after draining the queue so the writer
        # can't be mid-item (creating files inside the just-cleaned dir)
        # when rmtree runs.
        self._writer_busy = threading.Lock()
        self._writer_thread = threading.Thread(
            target=self._writer_loop,
            name="boundary-snapshot-writer",
            daemon=True,
        )
        self._writer_thread.start()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def save(
        self,
        request_id: str,
        token_count: int,
        snapshot_cache: list[Any],
        extract_cache_states_fn: Callable,
        *,
        block_size: int | None = None,
    ) -> bool:
        """Serialize snapshot to SSD (non-blocking).

        Must be called from the inference thread (Metal-safe for mx.eval).

        Parameters
        ----------
        request_id : str
            Unique request identifier.
        token_count : int
            Token boundary count.
        snapshot_cache : list
            Per-layer cache objects (None for skipped sliceable layers).
        extract_cache_states_fn : callable
            ``Scheduler._extract_cache_states`` — converts raw cache objects
            to ``List[Dict[str, Any]]``.
        block_size : int, optional
            When provided, compact append-only PoolingCache state to the
            current block's delta before buffering or writing it.

        Returns
        -------
        bool
            True if successfully enqueued for writing.
        """
        if not HAS_MLX:
            return False

        try:
            # 1. Extract dict-format states on inference thread.
            extracted, model_cache_config = extract_cache_states_fn(snapshot_cache)
            if not extracted:
                return False
            if block_size is not None:
                compact_pooling_cache_snapshot(extracted, token_count, block_size)

            # 2. Flatten tensors + metadata for safetensors serialization.
            tensors_raw, metadata = self._serialize_extracted(
                extracted, request_id, token_count
            )

            # 3. Buffer in pending writes for instant read-back.
            pw_key = (request_id, token_count)
            with self._pending_lock:
                self._pending_writes[pw_key] = {
                    "tensors_raw": tensors_raw,
                    "metadata": metadata,
                    "extracted": extracted,  # keep for cheap read-back
                }

            # 4. Compute file path and register.
            file_path = self._file_path(request_id, token_count)
            with self._registry_lock:
                self._file_registry.setdefault(request_id, {})[token_count] = file_path

            # 5. Enqueue for background write.
            try:
                self._write_queue.put_nowait((pw_key, tensors_raw, metadata, file_path))
            except queue.Full:
                # Roll back the pending + registry entries: with no
                # queue item the writer can never decrement
                # _cancelled_requests for this entry, so if a later
                # cleanup_request counts it the rid stays pinned in
                # _cancelled_requests forever and every subsequent
                # save under that rid is silently discarded by the
                # _is_cancelled gates. The previous "stays in memory
                # only" promise was already broken because cleanup
                # discards the in-memory copy anyway.
                logger.warning(
                    "Boundary snapshot write queue full, dropping "
                    "snapshot %s/%d",
                    request_id,
                    token_count,
                )
                with self._pending_lock:
                    self._pending_writes.pop(pw_key, None)
                with self._registry_lock:
                    req_files = self._file_registry.get(request_id)
                    if req_files is not None:
                        req_files.pop(token_count, None)
                        if not req_files:
                            self._file_registry.pop(request_id, None)
                return False

            return True

        except Exception as e:
            logger.debug("Failed to save boundary snapshot: %s", e)
            return False

    def load(
        self,
        request_id: str,
        token_count: int,
    ) -> list[dict[str, Any]] | None:
        """Load a snapshot, returning extracted cache state dicts.

        Checks the in-memory pending-writes buffer first (zero I/O), then
        falls back to reading the safetensors file from disk.

        Returns
        -------
        list or None
            List of per-layer dicts matching ``_extract_cache_states`` output
            format, or None on failure.
        """
        pw_key = (request_id, token_count)

        # Fast path: still in pending writes buffer.
        with self._pending_lock:
            pending = self._pending_writes.get(pw_key)
            if pending is not None:
                extracted = pending.get("extracted")
                if extracted is not None:
                    return extracted

                # Fallback: reconstruct from raw bytes.
                tensors_raw = pending.get("tensors_raw")
                metadata = pending.get("metadata")
                if tensors_raw and metadata:
                    return self._deserialize(tensors_raw, metadata)

        # Slow path: read from disk.
        file_path = self._file_path(request_id, token_count)
        if not file_path.exists():
            return None

        try:
            data = mx.load(str(file_path), return_metadata=True)
            if isinstance(data, tuple) and len(data) == 2:
                arrays, metadata = data
            else:
                return None
            return self._reconstruct_from_safetensors(arrays, metadata)
        except Exception as e:
            logger.debug(
                "Failed to load boundary snapshot %s/%d: %s",
                request_id,
                token_count,
                e,
            )
            return None

    def has(self, request_id: str, token_count: int) -> bool:
        """Check if a snapshot exists (in memory or on disk)."""
        pw_key = (request_id, token_count)
        with self._pending_lock:
            if pw_key in self._pending_writes:
                return True
        with self._registry_lock:
            req_files = self._file_registry.get(request_id)
            if req_files and token_count in req_files:
                return True
        return False

    def cleanup_request(self, request_id: str) -> None:
        """Delete all snapshot files and pending writes for a request.

        Caller must guarantee no async store_cache worker is still reading
        snapshots for this request — concurrent ``rmtree`` here would race
        the worker's :meth:`load` calls and silently strip block storage.
        :class:`omlx.scheduler.Scheduler` defers this call until the
        ``store_future`` for ``request_id`` is done.

        Acquires ``_writer_busy`` after marking the request cancelled so
        the writer thread can finish any item it is mid-processing first.
        Without this barrier the writer can pull an item, ``mkdir`` the
        request directory, write its temp file, then ``os.rename`` it
        into the final path *after* we have rmtree'd — leaving an
        orphaned file behind. The ``_cancelled_requests`` counter (held
        under ``_cancelled_lock``) catches the late-rename case if
        ``_writer_busy.acquire`` times out.

        Bounded with a timeout so a stuck I/O on the writer thread
        cannot deadlock request abort paths (called from scheduler's
        hot path at ~3 sites).

        The cancelled-counter is bumped additively and only when at
        least one pending item exists for the rid — see the inline
        comment at the bump site for the two distinct bugs that
        rules out (stale ``rid: 0`` after a timeout for an empty
        cleanup, and overwrites racing with re-entrant cleanup_request
        calls for the same rid).
        """
        # Atomically: count pending items for this rid, drop them, mark
        # the rid cancelled. Holding both locks during the snapshot is
        # required to keep the counter consistent with what the writer
        # will see — a save() call from another thread cannot interleave
        # an enqueue between our count and our cancellation mark.
        #
        # The bump is additive (``get + count``) and skipped entirely
        # when ``count == 0``. Both rules close real bugs:
        #   * Skip-on-zero: cleanup_request("X") for an rid with no
        #     pending items previously wrote ``cancelled[X] = 0`` then
        #     popped it on the acquired path. On the timeout fallback
        #     the pop never runs and the ``X: 0`` entry lingers for
        #     the process lifetime — every subsequent save() under
        #     that rid (or any later reuse of the same string) is
        #     discarded by the writer's ``_is_cancelled`` gates,
        #     which check key membership not value > 0. The counter
        #     must only exist when there is at least one in-flight
        #     item to drain it.
        #   * Additive: a re-entrant cleanup_request("X") for an rid
        #     that already has an in-flight cancellation must NOT
        #     overwrite the previous count. The writer's
        #     ``cleared_by_cleanup`` branch + ``_writer_busy`` lock
        #     together close the file-write race today, but the
        #     per-item dec_cancelled bookkeeping still has to balance.
        #     Overwriting drops the remaining decs on the floor; on
        #     the next ``save()`` under the same rid the writer would
        #     see a non-zero counter from the earlier batch and
        #     silently discard the new item.
        with self._pending_lock:
            keys_to_remove = [k for k in self._pending_writes if k[0] == request_id]
            count = len(keys_to_remove)
            for key in keys_to_remove:
                del self._pending_writes[key]
            if count > 0:
                with self._cancelled_lock:
                    self._cancelled_requests[request_id] = (
                        self._cancelled_requests.get(request_id, 0) + count
                    )

        # Remove from registry.
        with self._registry_lock:
            self._file_registry.pop(request_id, None)

        # Wait briefly for the writer to finish any item it had already
        # pulled. If it's genuinely stuck (slow disk, dead thread) fall
        # back to the cancelled-counter rescue rather than blocking the
        # caller.
        acquired = self._writer_busy.acquire(
            timeout=self._CLEANUP_REQUEST_TIMEOUT_S
        )
        try:
            # Remove files.
            req_dir = self._snapshot_dir / request_id
            if req_dir.exists():
                try:
                    shutil.rmtree(req_dir)
                except Exception as e:
                    logger.debug(
                        "Failed to clean up snapshots for %s: %s", request_id, e
                    )
        finally:
            if acquired:
                self._writer_busy.release()
                # Counter entry has done its job — we own the lock so all
                # _is_cancelled-gated work has either run or skipped. Drop
                # the counter so a future racing save() can't leave it
                # elevated forever. CRITICAL: only pop on the acquired
                # path. On timeout the writer is still mid-item and may
                # not yet have consulted ``_is_cancelled``; popping here
                # would defeat the late-rename rescue that the docstring
                # advertises as the timeout-fallback safety net.
                with self._cancelled_lock:
                    self._cancelled_requests.pop(request_id, None)
            else:
                logger.warning(
                    "cleanup_request(%s): writer thread did not yield "
                    "within %.1fs; relying on cancelled-counter rescue "
                    "for late-rename safety",
                    request_id,
                    self._CLEANUP_REQUEST_TIMEOUT_S,
                )

    def cleanup_all(self) -> None:
        """Delete all snapshot files for this store session.

        Synchronizes with the background writer: we drain the queue to
        prevent it from starting a new item, then acquire ``_writer_busy``
        to wait until any item it had already pulled finishes. Without
        this barrier the writer can create ``req-X/temp.safetensors``
        and ``os.rename`` it to its final path *after* we've already
        rmtree'd and recreated the snapshot directory, leaving an
        orphaned file behind.

        Threading: concurrent ``save()`` is safe because the writer
        consults ``_pending_writes`` and ``_is_cancelled`` while
        holding ``_writer_busy``, and ``cleanup_all`` clears both
        under the same lock before rmtree. The earlier "must run on
        the save() thread" constraint is therefore no longer required.
        """
        # Drain write queue so the writer thread doesn't process stale
        # items after the directory is deleted. Put_nowait the sentinel
        # back so shutdown still sees it; on Full just drop and let
        # shutdown re-issue.
        while True:
            try:
                item = self._write_queue.get_nowait()
                if item is None:  # Sentinel — put it back for shutdown.
                    try:
                        self._write_queue.put_nowait(item)
                    except queue.Full:
                        # Drop the sentinel; shutdown will re-enqueue.
                        # If cleanup_all is the LAST call before process
                        # exit without an explicit shutdown(), the writer
                        # thread will only be reaped on daemon teardown.
                        logger.debug(
                            "cleanup_all: dropped writer-sentinel on Full"
                        )
                    break
            except queue.Empty:
                break

        # Wait for the writer to finish any item it had already pulled.
        # When we own _writer_busy the writer is between items, and we
        # just drained the queue so no new item can start. Bounded so a
        # stuck writer (slow disk, dead thread) cannot deadlock callers
        # — scheduler calls cleanup_all() from its abort / reset hot
        # path. After the timeout we proceed anyway: the worst case is
        # an orphaned file in the recreated directory, which next
        # startup's cleanup_all() will clear.
        acquired = self._writer_busy.acquire(
            timeout=self._CLEANUP_ALL_TIMEOUT_S
        )
        try:
            if not acquired:
                logger.warning(
                    "cleanup_all: writer thread did not yield within "
                    "%.1fs; proceeding with rmtree — late-rename may "
                    "orphan a file under the recreated snapshot dir "
                    "until next startup.",
                    self._CLEANUP_ALL_TIMEOUT_S,
                )
            with self._pending_lock:
                self._pending_writes.clear()
            with self._registry_lock:
                self._file_registry.clear()
            with self._cancelled_lock:
                # Only safe to clear when we own _writer_busy — otherwise
                # a writer mid-_dec_cancelled would race. On timeout we
                # leave the counter intact so the rescue path stays
                # effective for in-flight items.
                if acquired:
                    self._cancelled_requests.clear()

            if self._snapshot_dir.exists():
                try:
                    shutil.rmtree(self._snapshot_dir)
                except Exception as e:
                    logger.debug(
                        "Failed to clean up all boundary snapshots: %s", e
                    )
            self._snapshot_dir.mkdir(parents=True, exist_ok=True)
        finally:
            if acquired:
                self._writer_busy.release()

    def shutdown(self) -> None:
        """Stop background writer thread."""
        self._shutdown.set()
        try:
            self._write_queue.put_nowait(None)  # Sentinel
        except queue.Full:
            pass
        self._writer_thread.join(timeout=5.0)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _is_cancelled(self, request_id: str) -> bool:
        """Thread-safe check for cancellation."""
        with self._cancelled_lock:
            return request_id in self._cancelled_requests

    def _dec_cancelled(self, request_id: str) -> None:
        """Decrement cancelled counter under lock; remove entry when
        exhausted. Atomic read-modify-write closes the underflow race
        between two writer-thread iterations / cleanup_all clears."""
        with self._cancelled_lock:
            remaining = self._cancelled_requests.get(request_id, 0) - 1
            if remaining <= 0:
                self._cancelled_requests.pop(request_id, None)
            else:
                self._cancelled_requests[request_id] = remaining

    def _file_path(self, request_id: str, token_count: int) -> Path:
        return self._snapshot_dir / request_id / f"{token_count}.safetensors"

    def _writer_loop(self) -> None:
        """Background thread that writes safetensors files."""
        while not self._shutdown.is_set():
            item = None
            try:
                item = self._write_queue.get(timeout=1.0)
            except queue.Empty:
                continue

            if item is None:  # Sentinel
                break

            # Hold _writer_busy for the entire item's lifetime so
            # cleanup_all() can serialize with us — otherwise it can
            # rmtree the snapshot directory while we're mid-write and
            # we'd recreate ``req-X/`` underneath it, leaving an
            # orphaned file after the cleanup returns.
            try:
                with self._writer_busy:
                    self._process_write_item(item)
            finally:
                # ``item`` contains ``tensors_raw``. If the thread waits for
                # the next queue entry without clearing this local, the frame
                # can pin a whole boundary snapshot after pending_writes was
                # cleaned up.
                item = None

    def _process_write_item(self, item) -> None:
        """Process one (pw_key, tensors_raw, metadata, file_path) queue item.

        Extracted from ``_writer_loop`` so the busy-lock can wrap it
        cleanly. Called only on the writer thread.
        """
        pw_key, tensors_raw, metadata, file_path = item

        # If cleanup_all or cleanup_request cleared this key from
        # _pending_writes while the item was in the writer's local hand
        # (i.e. between ``get()`` and entering ``with _writer_busy``),
        # treat the write as cancelled. This closes the late-rename
        # window where cleanup runs entirely between the writer's pull
        # and its busy-lock acquisition.
        with self._pending_lock:
            cleared_by_cleanup = pw_key not in self._pending_writes
        if cleared_by_cleanup:
            # If a timed-out cleanup_request bumped ``_cancelled_requests``
            # before clearing pending_writes, this item is one of the N
            # the counter is waiting on. Without this decrement the
            # counter would never reach zero, leaving the rid pinned in
            # ``_cancelled_requests`` for the process lifetime and
            # causing every subsequent write under that rid (or any
            # later reuse of the same string) to be silently discarded.
            if self._is_cancelled(pw_key[0]):
                self._dec_cancelled(pw_key[0])
            return

        # Skip writes for cancelled/cleaned-up requests.
        if self._is_cancelled(pw_key[0]):
            with self._pending_lock:
                self._pending_writes.pop(pw_key, None)
            try:
                req_dir = file_path.parent
                if req_dir.exists():
                    shutil.rmtree(req_dir)
            except Exception:
                pass
            self._dec_cancelled(pw_key[0])
            return

        temp_path = None
        try:
            file_path.parent.mkdir(parents=True, exist_ok=True)
            temp_path = file_path.with_name(file_path.stem + "_tmp.safetensors")
            _write_safetensors_no_mx(str(temp_path), tensors_raw, metadata)

            # Request may have been cleaned up while serializing.
            if self._is_cancelled(pw_key[0]):
                try:
                    if temp_path.exists():
                        temp_path.unlink()
                except Exception:
                    pass
                with self._pending_lock:
                    self._pending_writes.pop(pw_key, None)
                self._dec_cancelled(pw_key[0])
                return

            os.rename(str(temp_path), str(file_path))

            # Cleanup may race with a queued write; remove any late file.
            if self._is_cancelled(pw_key[0]):
                try:
                    if file_path.exists():
                        file_path.unlink()
                except Exception:
                    pass
                req_dir = file_path.parent
                try:
                    if req_dir.exists():
                        shutil.rmtree(req_dir)
                except Exception:
                    pass
                self._dec_cancelled(pw_key[0])
        except Exception as e:
            logger.debug("Background snapshot write failed: %s", e)
            for p in (temp_path, file_path):
                try:
                    if p is not None and p.exists():
                        p.unlink()
                except Exception:
                    pass
            # Same bookkeeping invariant as the early-return path: if
            # cleanup_request bumped the counter and the failure was a
            # side-effect of that cleanup (e.g. its rmtree pulled the
            # parent dir out from under our temp write), we still owe
            # one decrement. The _is_cancelled rescue blocks above all
            # return before this except clause runs, so we cannot
            # double-decrement.
            if self._is_cancelled(pw_key[0]):
                self._dec_cancelled(pw_key[0])
        finally:
            # Remove extracted cache objects from pending writes to free
            # memory, but keep tensors_raw for read-back until file is on
            # disk.
            with self._pending_lock:
                pending = self._pending_writes.get(pw_key)
                if pending is not None:
                    pending.pop("extracted", None)
                # If file was written successfully, remove entirely.
                if file_path.exists():
                    self._pending_writes.pop(pw_key, None)

    def _serialize_extracted(
        self,
        extracted: list[dict[str, Any]],
        request_id: str,
        token_count: int,
    ) -> tuple[dict[str, tuple[bytes, str, list[int]]], dict[str, str]]:
        """Convert extracted cache states to tensors_raw + metadata.

        Must be called on the inference thread (for mx.eval / _extract_tensor_bytes).
        """
        arrays: dict[str, Any] = {}  # name -> mx.array
        layer_info: list[dict[str, str]] = []

        for i, layer_state in enumerate(extracted):
            class_name = layer_state.get("class_name", "KVCache")
            cache_type = layer_state.get("cache_type", "KVCache")
            meta_state = layer_state.get("meta_state", ())
            state = layer_state.get("state", ())

            info: dict[str, str] = {
                "class_name": class_name,
                "cache_type": cache_type,
                "meta_state": json.dumps(list(meta_state) if meta_state else []),
            }
            pooling_delta_ranges = layer_state.get("pooling_delta_ranges")
            if pooling_delta_ranges:
                info["pooling_delta_ranges"] = json.dumps(pooling_delta_ranges)

            if (
                isinstance(state, list)
                and len(state) >= 1
                and all(isinstance(s, (list, tuple)) for s in state)
            ):
                # CacheList layer: ``state`` is a list of nested sub-state
                # tuples (one per sub-cache, e.g. RotatingKVCache +
                # PoolingCache for DeepSeek V4). Flatten as
                # ``layer_{i}_sub_{j}_state_{k}`` keys so reconstruction
                # can rebuild the nested shape. Mirror the flat branch's
                # has_tensors gate: a CacheList whose subs hold no tensor
                # at all (empty KV + untouched conv slots) carries no
                # restorable state.
                has_tensors = any(
                    hasattr(elem, "shape") for sub in state for elem in sub
                )
                if not has_tensors:
                    info["has_state"] = "false"
                    layer_info.append(info)
                    continue
                info["has_state"] = "true"
                info["sub_count"] = str(len(state))
                for j, sub_state in enumerate(state):
                    info[f"sub_{j}_count"] = str(len(sub_state))
                    for k, elem in enumerate(sub_state):
                        if not hasattr(elem, "shape"):
                            info[f"sub_{j}_missing_{k}"] = "1"
                            continue
                        if _has_zero_dim(elem):
                            arrays[f"layer_{i}_sub_{j}_state_{k}"] = mx.zeros((1,))
                            info[f"sub_{j}_zero_dim_{k}"] = _encode_shape(elem.shape)
                        else:
                            arrays[f"layer_{i}_sub_{j}_state_{k}"] = elem
            elif isinstance(state, (list, tuple)) and len(state) >= 1:
                # Flat N-tuple state (KVCache, RotatingKVCache, PoolingCache,
                # BatchKVCache). Store every element under
                # ``layer_{i}_state_{k}`` regardless of tuple length.
                has_tensors = any(hasattr(elem, "shape") for elem in state)
                if has_tensors:
                    info["has_state"] = "true"
                    info["state_count"] = str(len(state))
                    for k, elem in enumerate(state):
                        if not hasattr(elem, "shape"):
                            # Non-tensor element (None, scalar). Mark it so
                            # _deserialize can restore the gap.
                            info[f"missing_{k}"] = "1"
                            continue
                        if _has_zero_dim(elem):
                            arrays[f"layer_{i}_state_{k}"] = mx.zeros((1,))
                            info[f"zero_dim_{k}"] = _encode_shape(elem.shape)
                        else:
                            arrays[f"layer_{i}_state_{k}"] = elem
                else:
                    info["has_state"] = "false"
            else:
                info["has_state"] = "false"

            layer_info.append(info)

        # Materialize lazy tensors on inference thread.
        if arrays:
            mx.eval(*arrays.values())

        # Extract raw bytes (Metal-safe memoryview copy).
        tensors_raw = {}
        for name, arr in arrays.items():
            tensors_raw[name] = _extract_tensor_bytes(arr)

        metadata = {
            "request_id": request_id,
            "token_count": str(token_count),
            "num_layers": str(len(extracted)),
            "layer_info": json.dumps(layer_info),
        }

        return tensors_raw, metadata

    def _deserialize(
        self,
        tensors_raw: dict[str, tuple[bytes, str, list[int]]],
        metadata: dict[str, str],
    ) -> list[dict[str, Any]] | None:
        """Reconstruct extracted cache states from raw bytes + metadata."""
        try:
            num_layers = int(metadata["num_layers"])
            layer_info = json.loads(metadata["layer_info"])
        except (KeyError, ValueError, json.JSONDecodeError):
            return None

        result: list[dict[str, Any]] = []
        for i in range(num_layers):
            info = layer_info[i] if i < len(layer_info) else {}
            class_name = info.get("class_name", "KVCache")
            cache_type = info.get("cache_type", "KVCache")
            meta_state_json = info.get("meta_state", "[]")
            try:
                meta_state = tuple(json.loads(meta_state_json))
            except (ValueError, json.JSONDecodeError):
                meta_state = ()

            if info.get("has_state") == "true":
                # V3 path: state_count meta + layer_{i}_state_{k} keys.
                # V2 fallback: legacy layer_{i}_0/1 + zero_dim_0/1 keys
                # for snapshots written before the N-tuple migration.
                state = self._read_state_tuple_raw(tensors_raw, info, i)
                result.append(
                    {
                        "state": state,
                        "meta_state": meta_state,
                        "class_name": class_name,
                        "cache_type": cache_type,
                    }
                )
            else:
                # Placeholder for skipped sliceable layers.
                result.append(
                    {
                        "state": (),
                        "meta_state": meta_state,
                        "class_name": class_name,
                        "cache_type": cache_type,
                    }
                )

            if "pooling_delta_ranges" in info:
                try:
                    result[-1]["pooling_delta_ranges"] = json.loads(
                        info["pooling_delta_ranges"]
                    )
                except (TypeError, json.JSONDecodeError):
                    return None

        return result

    def _read_state_tuple_raw(
        self,
        tensors_raw: dict[str, tuple[bytes, str, list[int]]],
        info: dict[str, str],
        layer_idx: int,
    ) -> Any:
        """Read state for one layer from raw tensor bytes.

        Returns:
            - ``list`` of nested sub-state tuples for CacheList layers
              (``sub_count`` in info), or
            - ``tuple`` of N elements for flat layers (``state_count`` in
              info, V3 layout), or
            - 2-tuple from V2 polyfill (``layer_{i}_0`` / ``layer_{i}_1``).

        Missing elements come back as ``None``.
        """
        if "sub_count" in info:
            try:
                sub_count = int(info["sub_count"])
            except (ValueError, TypeError):
                return []
            sub_states: list[tuple[Any, ...]] = []
            for j in range(sub_count):
                count_key = f"sub_{j}_count"
                try:
                    count = int(info.get(count_key, "0"))
                except (ValueError, TypeError):
                    count = 0
                sub_elements: list[Any] = []
                for k in range(count):
                    if info.get(f"sub_{j}_missing_{k}") == "1":
                        sub_elements.append(None)
                        continue
                    key = f"layer_{layer_idx}_sub_{j}_state_{k}"
                    if key not in tensors_raw:
                        sub_elements.append(None)
                        continue
                    raw, dtype_str, shape = tensors_raw[key]
                    zd_marker = f"sub_{j}_zero_dim_{k}"
                    if zd_marker in info:
                        zd_shape = tuple(int(d) for d in info[zd_marker].split(","))
                        restored = _restore_tensor_from_bytes(raw, dtype_str, [1])
                        sub_elements.append(mx.zeros(zd_shape, dtype=restored.dtype))
                    else:
                        sub_elements.append(
                            _restore_tensor_from_bytes(raw, dtype_str, shape)
                        )
                sub_states.append(tuple(sub_elements))
            return sub_states

        if "state_count" in info:
            try:
                count = int(info["state_count"])
            except (ValueError, TypeError):
                return ()
            elements: list[Any] = []
            for k in range(count):
                if info.get(f"missing_{k}") == "1":
                    elements.append(None)
                    continue
                key = f"layer_{layer_idx}_state_{k}"
                if key not in tensors_raw:
                    elements.append(None)
                    continue
                raw, dtype_str, shape = tensors_raw[key]
                zd_marker = f"zero_dim_{k}"
                if zd_marker in info:
                    zd_shape = tuple(int(d) for d in info[zd_marker].split(","))
                    restored = _restore_tensor_from_bytes(raw, dtype_str, [1])
                    elements.append(mx.zeros(zd_shape, dtype=restored.dtype))
                else:
                    elements.append(_restore_tensor_from_bytes(raw, dtype_str, shape))
            return tuple(elements)

        # V2 polyfill — legacy 2-tuple snapshot.
        first = None
        second = None
        key_0 = f"layer_{layer_idx}_0"
        key_1 = f"layer_{layer_idx}_1"
        if key_0 in tensors_raw:
            raw, dtype_str, shape = tensors_raw[key_0]
            if "zero_dim_0" in info:
                zd_shape = tuple(int(d) for d in info["zero_dim_0"].split(","))
                first = _restore_tensor_from_bytes(raw, dtype_str, [1])
                first = mx.zeros(zd_shape, dtype=first.dtype)
            else:
                first = _restore_tensor_from_bytes(raw, dtype_str, shape)
        if key_1 in tensors_raw:
            raw, dtype_str, shape = tensors_raw[key_1]
            if "zero_dim_1" in info:
                zd_shape = tuple(int(d) for d in info["zero_dim_1"].split(","))
                second = _restore_tensor_from_bytes(raw, dtype_str, [1])
                second = mx.zeros(zd_shape, dtype=second.dtype)
            else:
                second = _restore_tensor_from_bytes(raw, dtype_str, shape)
        return (first, second) if first is not None else ()

    def _reconstruct_from_safetensors(
        self,
        arrays: dict[str, Any],
        metadata: dict[str, str],
    ) -> list[dict[str, Any]] | None:
        """Reconstruct from mx.load() result (arrays dict + metadata)."""
        try:
            num_layers = int(metadata["num_layers"])
            layer_info = json.loads(metadata["layer_info"])
        except (KeyError, ValueError, json.JSONDecodeError):
            return None

        result: list[dict[str, Any]] = []
        for i in range(num_layers):
            info = layer_info[i] if i < len(layer_info) else {}
            class_name = info.get("class_name", "KVCache")
            cache_type = info.get("cache_type", "KVCache")
            meta_state_json = info.get("meta_state", "[]")
            try:
                meta_state = tuple(json.loads(meta_state_json))
            except (ValueError, json.JSONDecodeError):
                meta_state = ()

            if info.get("has_state") == "true":
                state = self._read_state_tuple_arrays(arrays, info, i)
                result.append(
                    {
                        "state": state,
                        "meta_state": meta_state,
                        "class_name": class_name,
                        "cache_type": cache_type,
                    }
                )
            else:
                result.append(
                    {
                        "state": (),
                        "meta_state": meta_state,
                        "class_name": class_name,
                        "cache_type": cache_type,
                    }
                )

            if "pooling_delta_ranges" in info:
                try:
                    result[-1]["pooling_delta_ranges"] = json.loads(
                        info["pooling_delta_ranges"]
                    )
                except (TypeError, json.JSONDecodeError):
                    return None

        return result

    def _read_state_tuple_arrays(
        self,
        arrays: dict[str, Any],
        info: dict[str, str],
        layer_idx: int,
    ) -> Any:
        """N-tuple aware safetensors-loaded variant of
        ``_read_state_tuple_raw`` — sources tensors from a pre-decoded
        ``mx.array`` dict instead of raw bytes. Returns a list of nested
        tuples for CacheList layers (``sub_count`` in info) or a flat
        tuple otherwise.
        """
        if "sub_count" in info:
            try:
                sub_count = int(info["sub_count"])
            except (ValueError, TypeError):
                return []
            sub_states: list[tuple[Any, ...]] = []
            for j in range(sub_count):
                count_key = f"sub_{j}_count"
                try:
                    count = int(info.get(count_key, "0"))
                except (ValueError, TypeError):
                    count = 0
                sub_elements: list[Any] = []
                for k in range(count):
                    if info.get(f"sub_{j}_missing_{k}") == "1":
                        sub_elements.append(None)
                        continue
                    key = f"layer_{layer_idx}_sub_{j}_state_{k}"
                    tensor = arrays.get(key)
                    if tensor is None:
                        sub_elements.append(None)
                        continue
                    zd_marker = f"sub_{j}_zero_dim_{k}"
                    if zd_marker in info:
                        zd_shape = tuple(int(d) for d in info[zd_marker].split(","))
                        sub_elements.append(mx.zeros(zd_shape, dtype=tensor.dtype))
                    else:
                        sub_elements.append(tensor)
                sub_states.append(tuple(sub_elements))
            return sub_states

        if "state_count" in info:
            try:
                count = int(info["state_count"])
            except (ValueError, TypeError):
                return ()
            elements: list[Any] = []
            for k in range(count):
                if info.get(f"missing_{k}") == "1":
                    elements.append(None)
                    continue
                key = f"layer_{layer_idx}_state_{k}"
                tensor = arrays.get(key)
                if tensor is None:
                    elements.append(None)
                    continue
                zd_marker = f"zero_dim_{k}"
                if zd_marker in info:
                    zd_shape = tuple(int(d) for d in info[zd_marker].split(","))
                    elements.append(mx.zeros(zd_shape, dtype=tensor.dtype))
                else:
                    elements.append(tensor)
            return tuple(elements)

        # V2 polyfill.
        first = arrays.get(f"layer_{layer_idx}_0")
        second = arrays.get(f"layer_{layer_idx}_1")
        if "zero_dim_0" in info and first is not None:
            zd_shape = tuple(int(d) for d in info["zero_dim_0"].split(","))
            first = mx.zeros(zd_shape, dtype=first.dtype)
        if "zero_dim_1" in info and second is not None:
            zd_shape = tuple(int(d) for d in info["zero_dim_1"].split(","))
            second = mx.zeros(zd_shape, dtype=second.dtype)
        return (first, second) if first is not None else ()


# ===== FILE: omlx/cache/factory.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Factory for creating cache instances from configuration.

This module provides a unified way to instantiate cache components
based on configuration settings.

Note: oMLX only supports paged SSD-based caching. Memory KV cache is managed
by mlx-lm's BatchGenerator. When paged SSD cache is disabled, no oMLX caching
is performed.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from .paged_cache import PagedCacheManager
    from .prefix_cache import BlockAwarePrefixCache
    from .paged_ssd_cache import PagedSSDCacheManager
    from ..memory_monitor import MemoryMonitor


@dataclass
class CacheConfig:
    """
    Configuration for cache system (paged SSD-only).

    Note: oMLX only supports paged SSD-based caching via BlockAwarePrefixCache.
    Memory KV cache is managed entirely by mlx-lm's BatchGenerator.

    Attributes:
        block_size: Number of tokens per cache block.
        max_num_blocks: Maximum number of cache blocks.
        initial_blocks: Initial number of blocks to allocate.
        paged_ssd_cache_dir: Directory for paged SSD cache storage. If None, caching is disabled.
        max_paged_ssd_cache_size: Maximum size of paged SSD cache in bytes.
        max_kv_cache_memory: Maximum GPU memory for KV cache.
        model_name: Model name for cache isolation.
    """

    block_size: int = 64
    max_num_blocks: int = 1024
    initial_blocks: int = 256
    paged_ssd_cache_dir: Optional[Path] = None
    max_paged_ssd_cache_size: int = 100 * 1024 * 1024 * 1024  # 100GB
    max_kv_cache_memory: Optional[int] = None
    model_name: str = ""


class CacheFactory:
    """
    Factory for creating cache instances (paged SSD-only).

    This class provides static methods to create cache components
    based on a unified CacheConfig. oMLX only supports paged SSD-based caching
    via BlockAwarePrefixCache.

    Example:
        config = CacheConfig(
            block_size=64,
            paged_ssd_cache_dir=Path("/tmp/cache"),
        )

        paged_cache = CacheFactory.create_paged_cache(config)
        paged_ssd_cache = CacheFactory.create_paged_ssd_cache(config, "llama-3b")
        prefix_cache = CacheFactory.create_prefix_cache(
            config, paged_cache, paged_ssd_cache
        )
    """

    @staticmethod
    def create_paged_cache(
        config: CacheConfig,
        num_layers: Optional[int] = None,
    ) -> Optional["PagedCacheManager"]:
        """
        Create a PagedCacheManager instance.

        Args:
            config: Cache configuration.
            num_layers: Number of model layers (unused, kept for API compat).

        Returns:
            Configured PagedCacheManager instance, or None if paged SSD cache is disabled.
        """
        if config.paged_ssd_cache_dir is None:
            return None

        from .paged_cache import PagedCacheManager

        return PagedCacheManager(
            block_size=config.block_size,
            max_blocks=config.max_num_blocks,
            enable_caching=True,
            model_name=config.model_name,
            initial_blocks=config.initial_blocks,
        )

    @staticmethod
    def create_paged_ssd_cache(
        config: CacheConfig,
        model_name: Optional[str] = None,
    ) -> Optional["PagedSSDCacheManager"]:
        """
        Create a PagedSSDCacheManager instance.

        Args:
            config: Cache configuration.
            model_name: Override model name for cache isolation.

        Returns:
            Configured PagedSSDCacheManager instance, or None if disabled.
        """
        if config.paged_ssd_cache_dir is None:
            return None

        from .paged_ssd_cache import PagedSSDCacheManager

        cache_dir = config.paged_ssd_cache_dir
        if model_name:
            cache_dir = cache_dir / model_name

        return PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=config.max_paged_ssd_cache_size,
            expected_model_name=config.model_name or "",
            expected_block_size=config.block_size,
            expected_block_size_tokens=config.block_size,
            # No memory monitor available here, so leave
            # ``expected_kv_bytes_per_token`` at its 200 KB default.
            # The Scheduler construction path is the one that benefits
            # from the model-derived value; this factory is used mostly
            # by direct/test callers.
        )

    @staticmethod
    def create_prefix_cache(
        config: CacheConfig,
        model: Any = None,
        paged_cache: Optional["PagedCacheManager"] = None,
        paged_ssd_cache: Optional["PagedSSDCacheManager"] = None,
    ) -> Optional["BlockAwarePrefixCache"]:
        """
        Create a BlockAwarePrefixCache instance.

        Note: oMLX only supports paged SSD-based caching via BlockAwarePrefixCache.
        Returns None if paged SSD cache is disabled.

        Args:
            config: Cache configuration.
            model: Model instance for cache identification.
            paged_cache: PagedCacheManager for block-based caching.
            paged_ssd_cache: PagedSSDCacheManager for SSD storage.

        Returns:
            Configured BlockAwarePrefixCache instance, or None if disabled.
        """
        if config.paged_ssd_cache_dir is None or paged_cache is None:
            return None

        from .prefix_cache import BlockAwarePrefixCache

        return BlockAwarePrefixCache(
            model=model,
            paged_cache_manager=paged_cache,
            paged_ssd_cache_manager=paged_ssd_cache,
        )

    @staticmethod
    def create_memory_monitor(
        config: CacheConfig,
        paged_cache: Optional["PagedCacheManager"] = None,
    ) -> "MemoryMonitor":
        """
        Create a MemoryMonitor instance.

        Args:
            config: Cache configuration.
            paged_cache: PagedCacheManager to monitor.

        Returns:
            Configured MemoryMonitor instance.
        """
        from ..memory_monitor import MemoryMonitor

        # Use configured max KV cache memory, or default to 4GB
        max_kv_cache_memory = config.max_kv_cache_memory or (4 * 1024 * 1024 * 1024)

        monitor = MemoryMonitor(
            max_kv_cache_memory=max_kv_cache_memory,
        )

        if paged_cache is not None:
            monitor.set_paged_cache_manager(paged_cache)

        return monitor

    @staticmethod
    def create_full_cache_stack(
        config: CacheConfig,
        model: Any = None,
        num_layers: Optional[int] = None,
    ) -> dict:
        """
        Create a complete cache stack with all components.

        This is a convenience method that creates all cache components
        and wires them together. Returns all None values if paged SSD cache
        is disabled (oMLX only supports paged SSD-based caching).

        Args:
            config: Cache configuration.
            model: Model instance for cache identification.
            num_layers: Number of model layers.

        Returns:
            Dictionary with keys: paged_cache, paged_ssd_cache, prefix_cache, memory_monitor
        """
        paged_cache = None
        paged_ssd_cache = None
        prefix_cache = None
        memory_monitor = None

        # Only create cache components if paged SSD cache is enabled
        if config.paged_ssd_cache_dir is not None:
            paged_cache = CacheFactory.create_paged_cache(config, num_layers)
            paged_ssd_cache = CacheFactory.create_paged_ssd_cache(
                config, config.model_name
            )

            if paged_cache is not None and paged_ssd_cache is not None:
                paged_cache.set_paged_ssd_cache_manager(paged_ssd_cache)

            prefix_cache = CacheFactory.create_prefix_cache(
                config, model, paged_cache, paged_ssd_cache
            )

            if paged_cache is not None:
                memory_monitor = CacheFactory.create_memory_monitor(config, paged_cache)

        return {
            "paged_cache": paged_cache,
            "paged_ssd_cache": paged_ssd_cache,
            "prefix_cache": prefix_cache,
            "memory_monitor": memory_monitor,
        }


# ===== FILE: omlx/cache/hybrid_cache.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Hybrid cache configuration for models with mixed cache types.

This module provides configuration classes for models that use different
cache types across layers (e.g., Qwen3-Next with ArraysCache + KVCache).
"""

from dataclasses import dataclass, field
from typing import Any, List, Optional, Tuple
import logging

from .type_handlers import CacheType, CacheTypeHandler
from .type_registry import CacheTypeRegistry

logger = logging.getLogger(__name__)


@dataclass
class LayerCacheConfig:
    """Configuration for a single layer's cache.

    Attributes:
        layer_idx: Index of the layer
        cache_type: Type of cache used by this layer
        supports_block_slicing: Whether this layer's cache can be block-sliced
        class_name: Original mlx-lm class name
    """

    layer_idx: int
    cache_type: CacheType
    supports_block_slicing: bool
    class_name: str = ""

    @property
    def handler(self) -> CacheTypeHandler:
        """Get the handler for this layer's cache type."""
        return CacheTypeRegistry.get_handler(self.cache_type)


@dataclass
class ModelCacheConfig:
    """Cache configuration for an entire model.

    Contains per-layer cache type information and model-level metadata.
    This is used to properly handle hybrid models that mix different
    cache types (e.g., KVCache + ArraysCache).

    Attributes:
        model_name: Name of the model
        num_layers: Total number of layers
        layer_configs: Per-layer cache configurations
        is_hybrid: True if model uses multiple cache types
        sliceable_layer_count: Number of layers that support block slicing
    """

    model_name: str = ""
    num_layers: int = 0
    layer_configs: List[LayerCacheConfig] = field(default_factory=list)
    is_hybrid: bool = False
    sliceable_layer_count: int = 0

    @classmethod
    def from_cache_list(
        cls,
        cache_list: List[Any],
        model_name: str = "",
    ) -> "ModelCacheConfig":
        """Create configuration from mlx-lm cache list.

        Args:
            cache_list: List of cache objects from model.make_cache()
            model_name: Optional model name for identification

        Returns:
            ModelCacheConfig with per-layer type information
        """
        if not cache_list:
            return cls(model_name=model_name)

        layer_configs = []
        cache_types_seen = set()
        sliceable_count = 0
        max_window_size = 0

        for idx, cache_obj in enumerate(cache_list):
            cache_type = CacheTypeRegistry.detect_cache_type(cache_obj)
            handler = CacheTypeRegistry.get_handler(cache_type)
            class_name = type(cache_obj).__name__

            cache_types_seen.add(cache_type)
            if handler.supports_block_slicing:
                sliceable_count += 1

            # Extract window_size from RotatingKVCache layers
            if cache_type == CacheType.ROTATING_KVCACHE:
                window_size = getattr(cache_obj, "max_size", 0)
                if window_size > max_window_size:
                    max_window_size = window_size

            # Extract window_size from CacheList sub-caches (e.g., RotatingKVCache inside)
            if cache_type == CacheType.CACHE_LIST:
                sub_caches = getattr(cache_obj, "caches", ())
                for sub_c in sub_caches:
                    sub_type = CacheTypeRegistry.detect_cache_type(sub_c)
                    if sub_type == CacheType.ROTATING_KVCACHE:
                        ws = getattr(sub_c, "max_size", 0)
                        if ws > max_window_size:
                            max_window_size = ws

            layer_configs.append(
                LayerCacheConfig(
                    layer_idx=idx,
                    cache_type=cache_type,
                    supports_block_slicing=handler.supports_block_slicing,
                    class_name=class_name,
                )
            )

        config = cls(
            model_name=model_name,
            num_layers=len(cache_list),
            layer_configs=layer_configs,
            is_hybrid=len(cache_types_seen) > 1,
            sliceable_layer_count=sliceable_count,
        )
        config._max_window_size = max_window_size
        return config

    @classmethod
    def from_type_list(
        cls,
        cache_types: List[str],
        model_name: str = "",
    ) -> "ModelCacheConfig":
        """Create configuration from list of type names.

        Useful for reconstructing config from serialized metadata.

        Args:
            cache_types: List of cache type names (e.g., ["KVCache", "ArraysCache"])
            model_name: Optional model name

        Returns:
            ModelCacheConfig
        """
        if not cache_types:
            return cls(model_name=model_name)

        layer_configs = []
        cache_types_seen = set()
        sliceable_count = 0

        for idx, type_name in enumerate(cache_types):
            handler = CacheTypeRegistry.get_handler_by_class_name(type_name)
            cache_type = handler.cache_type

            cache_types_seen.add(cache_type)
            if handler.supports_block_slicing:
                sliceable_count += 1

            layer_configs.append(
                LayerCacheConfig(
                    layer_idx=idx,
                    cache_type=cache_type,
                    supports_block_slicing=handler.supports_block_slicing,
                    class_name=type_name,
                )
            )

        return cls(
            model_name=model_name,
            num_layers=len(cache_types),
            layer_configs=layer_configs,
            is_hybrid=len(cache_types_seen) > 1,
            sliceable_layer_count=sliceable_count,
        )

    def get_sliceable_layers(self) -> List[int]:
        """Get indices of layers that support block slicing.

        Returns:
            List of layer indices
        """
        return [cfg.layer_idx for cfg in self.layer_configs if cfg.supports_block_slicing]

    def get_non_sliceable_layers(self) -> List[int]:
        """Get indices of layers that don't support block slicing.

        Returns:
            List of layer indices (e.g., ArraysCache, RotatingKVCache layers)
        """
        return [
            cfg.layer_idx for cfg in self.layer_configs if not cfg.supports_block_slicing
        ]

    def get_layer_type(self, layer_idx: int) -> CacheType:
        """Get cache type for a specific layer.

        Args:
            layer_idx: Layer index

        Returns:
            CacheType enum
        """
        if 0 <= layer_idx < len(self.layer_configs):
            return self.layer_configs[layer_idx].cache_type
        return CacheType.KVCACHE  # Default

    def get_layer_handler(self, layer_idx: int) -> CacheTypeHandler:
        """Get handler for a specific layer.

        Args:
            layer_idx: Layer index

        Returns:
            CacheTypeHandler instance
        """
        if 0 <= layer_idx < len(self.layer_configs):
            return self.layer_configs[layer_idx].handler
        return CacheTypeRegistry.get_handler(CacheType.KVCACHE)

    def get_type_names(self) -> List[str]:
        """Get list of cache type names for serialization.

        Returns:
            List of class name strings
        """
        return [cfg.class_name for cfg in self.layer_configs]

    def get_meta_states(self, cache_list: List[Any]) -> List[Tuple]:
        """Extract meta_states from cache objects.

        For CacheList layers, the meta_state is a composite:
        ([class_names], [sub_meta_states]).

        Args:
            cache_list: List of cache objects

        Returns:
            List of meta_state tuples
        """
        meta_states = []
        for idx, cache_obj in enumerate(cache_list):
            if idx < len(self.layer_configs):
                handler = self.layer_configs[idx].handler

                if self.layer_configs[idx].cache_type == CacheType.CACHE_LIST:
                    # CacheList: extract composite meta_state from sub-caches
                    state = handler.extract_state(cache_obj)
                    meta_state = (
                        state.get("sub_class_names", []),
                        state.get("sub_meta_states", []),
                    )
                    meta_states.append(meta_state)
                else:
                    state = handler.extract_state(cache_obj)
                    meta_state = state.get("meta_state", ())
                    if not meta_state:
                        # Try to extract offset at minimum
                        offset = state.get("offset", 0)
                        meta_state = (offset,)
                    meta_states.append(meta_state)
            else:
                meta_states.append(())
        return meta_states

    def supports_full_block_slicing(self) -> bool:
        """Check if all layers support block slicing.

        Returns:
            True if all layers can be block-sliced
        """
        return self.sliceable_layer_count == self.num_layers

    def has_rotating_layers(self) -> bool:
        """Check if model has any RotatingKVCache layers.

        Returns:
            True if any layer uses RotatingKVCache
        """
        return any(
            cfg.cache_type == CacheType.ROTATING_KVCACHE for cfg in self.layer_configs
        )

    def get_max_window_size(self) -> int:
        """Get maximum window size from RotatingKVCache layers.

        The window size is extracted from RotatingKVCache meta_state during
        from_cache_list() construction, or can be set manually via
        _max_window_size attribute.

        Returns:
            Maximum window_size, or 0 if no RotatingKVCache layers
        """
        return getattr(self, "_max_window_size", 0)

    def __repr__(self) -> str:
        types = [cfg.class_name for cfg in self.layer_configs[:5]]
        if len(self.layer_configs) > 5:
            types.append(f"... +{len(self.layer_configs) - 5} more")
        return (
            f"ModelCacheConfig(model='{self.model_name}', "
            f"layers={self.num_layers}, "
            f"hybrid={self.is_hybrid}, "
            f"types={types})"
        )


def create_default_kvcache_config(num_layers: int, model_name: str = "") -> ModelCacheConfig:
    """Create a default KVCache-only configuration.

    Convenience function for models that use only KVCache.

    Args:
        num_layers: Number of model layers
        model_name: Optional model name

    Returns:
        ModelCacheConfig with all KVCache layers
    """
    layer_configs = [
        LayerCacheConfig(
            layer_idx=idx,
            cache_type=CacheType.KVCACHE,
            supports_block_slicing=True,
            class_name="KVCache",
        )
        for idx in range(num_layers)
    ]

    return ModelCacheConfig(
        model_name=model_name,
        num_layers=num_layers,
        layer_configs=layer_configs,
        is_hybrid=False,
        sliceable_layer_count=num_layers,
    )


# ===== FILE: omlx/cache/interface.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Cache manager interface for oMLX.

This module defines the abstract interface that all cache implementations
should follow for consistency across the codebase.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Tuple

from .stats import BaseCacheStats


class CacheManager(ABC):
    """
    Abstract interface for all cache implementations.

    This interface provides a consistent API for:
    - Prefix cache (trie-based LRU)
    - Paged cache (block-based KV cache)
    - VLM cache (vision-language model cache)
    - Paged SSD cache (disk-based persistence)
    """

    @abstractmethod
    def fetch(self, key: Any) -> Tuple[Optional[Any], bool]:
        """
        Fetch a value from the cache.

        Args:
            key: The cache key (varies by implementation).

        Returns:
            Tuple of (value, hit) where hit is True if found.
        """
        pass

    @abstractmethod
    def store(self, key: Any, value: Any) -> bool:
        """
        Store a value in the cache.

        Args:
            key: The cache key.
            value: The value to store.

        Returns:
            True if stored successfully.
        """
        pass

    @abstractmethod
    def evict(self, key: Any) -> bool:
        """
        Evict a specific entry from the cache.

        Args:
            key: The cache key to evict.

        Returns:
            True if evicted, False if not found.
        """
        pass

    @abstractmethod
    def clear(self) -> int:
        """
        Clear all entries from the cache.

        Returns:
            Number of entries cleared.
        """
        pass

    @abstractmethod
    def get_stats(self) -> BaseCacheStats:
        """
        Get cache statistics.

        Returns:
            BaseCacheStats or subclass with cache metrics.
        """
        pass

    @property
    @abstractmethod
    def size(self) -> int:
        """
        Get the current number of entries in the cache.

        Returns:
            Number of cached entries.
        """
        pass

    @property
    @abstractmethod
    def max_size(self) -> int:
        """
        Get the maximum capacity of the cache.

        Returns:
            Maximum number of entries.
        """
        pass

    @property
    def utilization(self) -> float:
        """
        Get cache utilization as a fraction.

        Returns:
            Utilization between 0.0 and 1.0.
        """
        if self.max_size == 0:
            return 0.0
        return self.size / self.max_size


# ===== FILE: omlx/cache/observability.py =====

# SPDX-License-Identifier: Apache-2.0
import threading
import time
from collections import deque
from typing import Any


_DEFAULT_WINDOWS = (60, 300, 900)
_MAX_SNAPSHOTS = 90
_MIN_INTERVAL = 10.0


class CacheRateTracker:

    def __init__(
        self,
        max_snapshots: int = _MAX_SNAPSHOTS,
        min_interval: float = _MIN_INTERVAL,
    ):
        self._snapshots: deque[tuple[float, dict[str, int]]] = deque(
            maxlen=max_snapshots
        )
        self._min_interval = min_interval
        self._lock = threading.Lock()

    def maybe_snapshot(self, counters: dict[str, int]) -> bool:
        with self._lock:
            now = time.monotonic()
            if self._snapshots and (now - self._snapshots[-1][0]) < self._min_interval:
                return False
            self._snapshots.append((now, dict(counters)))
            return True

    def get_rates(
        self, windows: tuple[int, ...] = _DEFAULT_WINDOWS
    ) -> dict[str, Any]:
        with self._lock:
            if not self._snapshots:
                return {"windows": {}, "cumulative": {}}

            now = self._snapshots[-1][0]
            newest = self._snapshots[-1][1]

            window_rates = {}
            for w in windows:
                label = _window_label(w)
                baseline_ts = None
                baseline_counters = None
                for ts, counters in self._snapshots:
                    if (now - ts) <= w:
                        baseline_ts, baseline_counters = ts, counters
                        break
                if baseline_ts is None:
                    baseline_ts, baseline_counters = self._snapshots[0]
                elapsed = now - baseline_ts
                if elapsed < 1.0:
                    window_rates[label] = {}
                    continue
                window_rates[label] = _compute_window(
                    baseline_counters, newest, elapsed
                )

            cumulative = _compute_cumulative(newest)
            return {"windows": window_rates, "cumulative": cumulative}

    def snapshot_and_get_rates(
        self,
        counters: dict[str, int],
        windows: tuple[int, ...] = _DEFAULT_WINDOWS,
    ) -> dict[str, Any]:
        self.maybe_snapshot(counters)
        return self.get_rates(windows)

    def clear(self) -> None:
        with self._lock:
            self._snapshots.clear()


def _window_label(seconds: int) -> str:
    if seconds < 60:
        return f"{seconds}s"
    return f"{seconds // 60}m"


def _safe_ratio(numerator: int, denominator: int) -> float:
    if denominator == 0:
        return 0.0
    return numerator / denominator


def _compute_window(
    old: dict[str, int], new: dict[str, int], elapsed: float
) -> dict[str, Any]:
    def delta(key: str) -> int:
        return max(0, new.get(key, 0) - old.get(key, 0))

    d_prefix_hits = delta("prefix_hits")
    d_prefix_misses = delta("prefix_misses")
    d_evictions = delta("evictions")
    d_ssd_hot = delta("ssd_hot_hits")
    d_ssd_disk = delta("ssd_disk_loads")
    d_tokens_matched = delta("prefix_tokens_matched")
    d_tokens_requested = delta("prefix_tokens_requested")

    minutes = elapsed / 60.0

    return {
        "prefix_hit_rate": round(
            _safe_ratio(d_prefix_hits, d_prefix_hits + d_prefix_misses), 4
        ),
        "prefix_hits": d_prefix_hits,
        "prefix_misses": d_prefix_misses,
        "prefix_match_efficiency": round(
            _safe_ratio(d_tokens_matched, d_tokens_requested), 4
        ),
        "evictions": d_evictions,
        "eviction_rate_per_min": round(d_evictions / minutes, 2) if minutes > 0 else 0.0,
        "ssd_hot_hits": d_ssd_hot,
        "ssd_disk_loads": d_ssd_disk,
        "ssd_hot_rate": round(
            _safe_ratio(d_ssd_hot, d_ssd_hot + d_ssd_disk), 4
        ),
    }


def _compute_cumulative(counters: dict[str, int]) -> dict[str, Any]:
    prefix_hits = counters.get("prefix_hits", 0)
    prefix_misses = counters.get("prefix_misses", 0)
    ssd_hot = counters.get("ssd_hot_hits", 0)
    ssd_disk = counters.get("ssd_disk_loads", 0)
    tokens_matched = counters.get("prefix_tokens_matched", 0)
    tokens_requested = counters.get("prefix_tokens_requested", 0)

    return {
        "prefix_hits": prefix_hits,
        "prefix_misses": prefix_misses,
        "prefix_hit_rate": round(_safe_ratio(prefix_hits, prefix_hits + prefix_misses), 4),
        "prefix_tokens_saved": counters.get("prefix_tokens_saved", 0),
        "prefix_match_efficiency": round(
            _safe_ratio(tokens_matched, tokens_requested), 4
        ),
        "evictions": counters.get("evictions", 0),
        "ssd_hot_hits": ssd_hot,
        "ssd_disk_loads": ssd_disk,
        "ssd_saves": counters.get("ssd_saves", 0),
        "hot_cache_evictions": counters.get("hot_cache_evictions", 0),
        "hot_cache_promotions": counters.get("hot_cache_promotions", 0),
        "ssd_hot_rate": round(_safe_ratio(ssd_hot, ssd_hot + ssd_disk), 4),
    }


# ===== FILE: omlx/cache/paged_cache.py =====

# SPDX-License-Identifier: Apache-2.0
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
"""
Paged KV Cache Manager for oMLX.

This module implements block-based paged KV cache management following vLLM's
architecture (vllm/v1/core/block_pool.py), adapted for MLX on Apple Silicon.

Key components:
- KVCacheBlock: Metadata for each cache block with doubly linked list pointers
- FreeKVCacheBlockQueue: O(1) doubly linked list for LRU block allocation
- BlockHashToBlockMap: Hash-to-block cache for prefix caching
- PagedCacheManager: Main manager with block allocation, prefix caching, and COW

Features:
- Block-based allocation (configurable tokens per block)
- Reference counting for shared blocks
- Copy-on-Write (COW) for efficient prefix sharing
- LRU eviction using doubly linked list (O(1) operations)
- Chain hashing for prefix caching (hash depends on parent block)

Reference: vLLM v1 - vllm/v1/core/block_pool.py, vllm/v1/core/kv_cache_utils.py
"""

from __future__ import annotations

import hashlib
import logging
import threading
import time
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from typing import Any, Dict, List, NewType, Optional, Tuple

from .interface import CacheManager
from .stats import BaseCacheStats, PagedCacheStats

logger = logging.getLogger(__name__)

# Type alias for block hash (content-based hash for prefix caching)
BlockHash = NewType("BlockHash", bytes)


def resolve_block_extra_keys(
    block_end: int,
    extra_keys: Optional[Tuple[Any, ...]] = None,
    extra_key_token_start: Optional[int] = None,
    extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
) -> Optional[Tuple[Any, ...]]:
    """Resolve which cache key salt applies to a block ending at ``block_end``.

    ``extra_key_ranges`` takes precedence over ``extra_keys`` and is intended
    for segmented VLM cache keying.

    ``extra_key_ranges`` must be sorted ascending by start position. The
    function uses early break on the first non-matching entry, so unsorted
    input will produce incorrect results.
    """
    if extra_key_ranges:
        assert all(
            extra_key_ranges[i][0] <= extra_key_ranges[i + 1][0]
            for i in range(len(extra_key_ranges) - 1)
        ), "extra_key_ranges must be sorted ascending by start position"
        selected = None
        for start, keys in extra_key_ranges:
            if block_end > start:
                selected = keys
            else:
                break
        return selected
    if extra_keys is not None and (
        extra_key_token_start is None or block_end > extra_key_token_start
    ):
        return extra_keys
    return None


def compute_block_hash(
    parent_hash: Optional[BlockHash],
    token_ids: List[int],
    extra_keys: Optional[Tuple[Any, ...]] = None,
    model_name: Optional[str] = None,
) -> BlockHash:
    """
    Compute hash for a block based on its content and parent block.

    This enables prefix caching by creating a chain of hashes where
    each block's hash depends on all previous blocks (similar to vLLM).

    Args:
        parent_hash: Hash of the previous block, or None for first block
        token_ids: Token IDs in this block
        extra_keys: Additional keys (e.g., LoRA, multimodal)
        model_name: Model name for cache isolation between different models

    Returns:
        Content-based hash for this block
    """
    hasher = hashlib.sha256()

    # Include model name first to isolate caches between different models
    if model_name:
        hasher.update(model_name.encode("utf-8"))

    # Include parent hash for chain
    if parent_hash:
        hasher.update(parent_hash)
    else:
        # Use fixed seed for reproducibility
        hasher.update(b"omlx-root")

    # Include token content
    hasher.update(bytes(str(tuple(token_ids)), "utf-8"))

    # Include extra keys if present
    if extra_keys:
        hasher.update(bytes(str(extra_keys), "utf-8"))

    return BlockHash(hasher.digest())


# =============================================================================
# KVCacheBlock - Following vLLM's design
# =============================================================================

@dataclass
class CacheBlock:
    """
    KV cache block metadata following vLLM's design.

    Each block represents a fixed number of tokens (block_size) worth
    of KV cache data. Blocks can be shared across requests via
    reference counting for prefix caching.

    NOTE: In paged SSD-only mode, blocks do NOT store cache_data in GPU memory.
    All KV cache data is stored on paged SSD via PagedSSDCacheManager, and only
    loaded when needed for inference via BatchGenerator.

    Attributes:
        block_id: Physical block index (0 to num_blocks - 1)
        ref_count: Reference count for sharing (0 = can be evicted)
        block_hash: Content hash for prefix caching and paged SSD storage key
        prev_free_block: Previous block in free list (doubly linked)
        next_free_block: Next block in free list (doubly linked)
        is_null: True if this is the null/placeholder block
        token_count: Number of tokens stored in this block
    """

    block_id: int
    ref_count: int = 0
    block_hash: Optional[BlockHash] = None

    # Doubly linked list pointers for FreeKVCacheBlockQueue
    prev_free_block: Optional["CacheBlock"] = None
    next_free_block: Optional["CacheBlock"] = None

    # Special flags
    is_null: bool = False

    # Metadata
    token_count: int = 0
    last_access: float = field(default_factory=time.time)

    def is_full(self, block_size: int) -> bool:
        """Check if block is at capacity."""
        return self.token_count >= block_size

    def is_shared(self) -> bool:
        """Check if block is shared (ref_count > 1)."""
        return self.ref_count > 1

    def reset_hash(self) -> None:
        """Reset block hash when evicted from cache."""
        self.block_hash = None

    def touch(self) -> None:
        """Update last access time."""
        self.last_access = time.time()

    def __repr__(self) -> str:
        prev_id = self.prev_free_block.block_id if self.prev_free_block else None
        next_id = self.next_free_block.block_id if self.next_free_block else None
        hash_str = f", hash={self.block_hash.hex()[:8]}..." if self.block_hash else ""
        return (
            f"CacheBlock(id={self.block_id}, ref={self.ref_count}, "
            f"tokens={self.token_count}, prev={prev_id}, next={next_id}{hash_str})"
        )


# =============================================================================
# FreeKVCacheBlockQueue - O(1) Doubly Linked List (vLLM style)
# =============================================================================

class FreeKVCacheBlockQueue:
    """
    Doubly linked list of free blocks following vLLM's design.

    Provides O(1) operations for:
    - popleft(): Allocate block from front (LRU order)
    - remove(): Remove block from middle (when touched by cache hit)
    - append(): Return block to end (when freed)

    The queue maintains LRU eviction order:
    - Front = least recently used (evict first)
    - Back = most recently used (evict last)

    Uses fake head/tail sentinels to simplify edge cases.
    """

    def __init__(self, blocks: List[CacheBlock]) -> None:
        """
        Initialize queue with all blocks as free.

        Args:
            blocks: List of all CacheBlock objects
        """
        self.num_free_blocks = len(blocks)

        # Initialize doubly linked list
        for i in range(len(blocks)):
            if i > 0:
                blocks[i].prev_free_block = blocks[i - 1]
            if i < len(blocks) - 1:
                blocks[i].next_free_block = blocks[i + 1]

        # Create sentinel nodes (never popped)
        self.fake_head = CacheBlock(block_id=-1)
        self.fake_tail = CacheBlock(block_id=-2)

        if blocks:
            self.fake_head.next_free_block = blocks[0]
            blocks[0].prev_free_block = self.fake_head
            self.fake_tail.prev_free_block = blocks[-1]
            blocks[-1].next_free_block = self.fake_tail
        else:
            self.fake_head.next_free_block = self.fake_tail
            self.fake_tail.prev_free_block = self.fake_head

    def popleft(self) -> CacheBlock:
        """
        Pop and return the first (LRU) free block.

        Raises:
            ValueError: If no free blocks available
        """
        if self.fake_head.next_free_block is self.fake_tail:
            raise ValueError("No free blocks available")

        block = self.fake_head.next_free_block
        assert block is not None

        # Remove from list
        self.fake_head.next_free_block = block.next_free_block
        if block.next_free_block:
            block.next_free_block.prev_free_block = self.fake_head

        block.prev_free_block = None
        block.next_free_block = None
        self.num_free_blocks -= 1

        return block

    def popleft_n(self, n: int) -> List[CacheBlock]:
        """
        Pop n blocks from the front.

        Args:
            n: Number of blocks to allocate

        Returns:
            List of n free blocks

        Raises:
            AssertionError: If not enough free blocks
        """
        if n == 0:
            return []

        assert self.num_free_blocks >= n, f"Need {n} blocks, have {self.num_free_blocks}"

        result = []
        curr = self.fake_head.next_free_block

        for _ in range(n):
            assert curr is not None and curr is not self.fake_tail
            result.append(curr)
            last = curr
            curr = curr.next_free_block
            # Clear pointers
            last.prev_free_block = None
            last.next_free_block = None

        # Reconnect list
        self.fake_head.next_free_block = curr
        if curr:
            curr.prev_free_block = self.fake_head

        self.num_free_blocks -= n
        return result

    def remove(self, block: CacheBlock) -> None:
        """
        Remove a block from the middle of the queue.

        Used when a free block is "touched" (reused by prefix cache hit).

        Args:
            block: Block to remove

        Raises:
            RuntimeError: If block not in queue
        """
        if block.prev_free_block is None or block.next_free_block is None:
            raise RuntimeError(f"Block {block.block_id} not in free queue")

        # Unlink
        block.prev_free_block.next_free_block = block.next_free_block
        block.next_free_block.prev_free_block = block.prev_free_block
        block.prev_free_block = None
        block.next_free_block = None

        self.num_free_blocks -= 1

    def append(self, block: CacheBlock) -> None:
        """
        Append a block to the end (MRU position).

        Args:
            block: Block to append
        """
        last = self.fake_tail.prev_free_block
        assert last is not None

        last.next_free_block = block
        block.prev_free_block = last
        block.next_free_block = self.fake_tail
        self.fake_tail.prev_free_block = block

        self.num_free_blocks += 1

    def append_n(self, blocks: List[CacheBlock]) -> None:
        """
        Append multiple blocks to the end.

        Args:
            blocks: Blocks to append (in order)
        """
        if not blocks:
            return

        last = self.fake_tail.prev_free_block
        assert last is not None

        for block in blocks:
            block.prev_free_block = last
            last.next_free_block = block
            last = block

        last.next_free_block = self.fake_tail
        self.fake_tail.prev_free_block = last

        self.num_free_blocks += len(blocks)

    def get_all_free_blocks(self) -> List[CacheBlock]:
        """Get all free blocks (for testing)."""
        result = []
        curr = self.fake_head.next_free_block
        while curr and curr is not self.fake_tail:
            result.append(curr)
            curr = curr.next_free_block
        return result


# =============================================================================
# BlockHashToBlockMap - Hash-based prefix cache (vLLM style)
# =============================================================================

class BlockHashToBlockMap:
    """
    Cache mapping block hashes to blocks for prefix caching.

    Follows vLLM's design where the same hash can map to multiple
    blocks (for different KV cache groups in hybrid models).
    """

    def __init__(self) -> None:
        self._cache: Dict[BlockHash, CacheBlock | Dict[int, CacheBlock]] = {}

    def get_block(self, block_hash: BlockHash) -> Optional[CacheBlock]:
        """Get any block with the given hash."""
        blocks = self._cache.get(block_hash)
        if blocks is None:
            return None
        if isinstance(blocks, CacheBlock):
            return blocks
        if isinstance(blocks, dict):
            return next(iter(blocks.values()))
        return None

    def insert(self, block_hash: BlockHash, block: CacheBlock) -> None:
        """Insert a block into the cache."""
        existing = self._cache.get(block_hash)
        if existing is None:
            self._cache[block_hash] = block
        elif isinstance(existing, CacheBlock):
            self._cache[block_hash] = {
                existing.block_id: existing,
                block.block_id: block,
            }
        elif isinstance(existing, dict):
            existing[block.block_id] = block

    def pop(self, block_hash: BlockHash, block_id: int) -> Optional[CacheBlock]:
        """Remove and return a specific block from the cache."""
        blocks = self._cache.pop(block_hash, None)
        if blocks is None:
            return None

        if isinstance(blocks, CacheBlock):
            if blocks.block_id == block_id:
                return blocks
            # Wrong block ID, put it back
            self._cache[block_hash] = blocks
            return None

        if isinstance(blocks, dict):
            block = blocks.pop(block_id, None)
            if blocks:  # Still has other blocks
                self._cache[block_hash] = blocks
            return block

        return None

    def __len__(self) -> int:
        return len(self._cache)

    def clear(self) -> None:
        self._cache.clear()


# =============================================================================
# BlockTable - Per-request block mapping
# =============================================================================

@dataclass
class BlockTable:
    """
    Per-request block table mapping logical to physical blocks.

    Similar to vLLM's block table, this maps a request's token positions
    to physical cache blocks.

    Attributes:
        request_id: Unique request identifier
        block_ids: List of physical block IDs
        num_tokens: Total number of cached tokens
    """

    request_id: str
    block_ids: List[int] = field(default_factory=list)
    num_tokens: int = 0

    def add_block(self, block_id: int, num_tokens: int) -> None:
        """Add a block to the table."""
        self.block_ids.append(block_id)
        self.num_tokens += num_tokens

    def __len__(self) -> int:
        return len(self.block_ids)

    def copy(self, new_request_id: str) -> "BlockTable":
        """Create a copy with new request ID."""
        return BlockTable(
            request_id=new_request_id,
            block_ids=self.block_ids.copy(),
            num_tokens=self.num_tokens,
        )


# =============================================================================
# PagedCacheManager - Main manager (vLLM BlockPool style)
# =============================================================================

class PagedCacheManager(CacheManager):
    """
    Paged KV cache manager following vLLM's BlockPool architecture.

    Features:
    - Block allocation/deallocation with reference counting
    - Prefix sharing via chain-based hash deduplication
    - Copy-on-Write for efficient forking
    - O(1) LRU eviction using doubly linked list

    Implements the CacheManager ABC interface for consistency with other
    cache implementations in oMLX.

    Args:
        block_size: Number of tokens per block (default: 64)
        max_blocks: Maximum number of blocks to allocate (default: 1000)
        enable_caching: Whether to enable prefix caching (default: True)
    """

    def __init__(
        self,
        block_size: int = 64,
        max_blocks: int = 1000,
        enable_caching: bool = True,
        model_name: str = "",
        initial_blocks: int = 256,
    ):
        self.block_size = block_size
        self.max_blocks = max_blocks
        self.enable_caching = enable_caching
        self.model_name = model_name
        self.initial_blocks = initial_blocks

        # Warn if model_name is not set (cache isolation may not work)
        if not model_name:
            logger.warning(
                "PagedCacheManager initialized without model_name. "
                "Cache isolation between models may not work correctly."
            )

        # Dynamic allocation tracking
        # Only create initial blocks; grow dynamically as needed up to max_blocks
        initial_count = min(initial_blocks, max_blocks)
        self._current_allocated_count = initial_count

        # Create only initial blocks (memory optimization)
        self.blocks: List[CacheBlock] = [
            CacheBlock(block_id=i) for i in range(initial_count)
        ]

        # Free block queue (doubly linked list for O(1) LRU)
        self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)

        # Hash-to-block cache for prefix caching
        self.cached_block_hash_to_block = BlockHashToBlockMap()

        # Request to block table mapping
        self.request_tables: Dict[str, BlockTable] = {}

        # Allocated blocks (for fast lookup)
        self.allocated_blocks: Dict[int, CacheBlock] = {}

        # Reserve null block (block 0) - never freed
        self.null_block = self.free_block_queue.popleft()
        self.null_block.is_null = True
        self.null_block.ref_count = 1
        self.allocated_blocks[self.null_block.block_id] = self.null_block

        # Statistics - track actual created blocks, not max
        self.stats = PagedCacheStats(
            total_blocks=initial_count,
            allocated_blocks=1,  # null block
            free_blocks=initial_count - 1,
        )

        # Thread safety
        self._lock = threading.RLock()

        # paged SSD cache manager for storage (set via set_paged_ssd_cache_manager)
        self._paged_ssd_cache_manager: Optional[Any] = None

        # Lifecycle hooks for hash-keyed side indexes (e.g. the prefix index
        # in BlockAwarePrefixCache). on_block_hash_dropped fires when a
        # (hash -> block) association ceases to exist; on_hash_map_cleared
        # fires after a wholesale hash-map clear. Callbacks may run with
        # self._lock held and must not call back into this manager.
        self.on_block_hash_dropped: Callable[[BlockHash], None] | None = None
        self.on_hash_map_cleared: Callable[[], None] | None = None

        logger.info(
            f"PagedCacheManager initialized: block_size={block_size}, "
            f"initial_blocks={initial_count}, max_blocks={max_blocks}, "
            f"max_tokens={block_size * max_blocks}"
        )

    def set_paged_ssd_cache_manager(self, paged_ssd_cache_manager: Any) -> None:
        """
        Set paged SSD cache manager for tiered storage.

        When set, evicted blocks will be saved to paged SSD before clearing
        GPU memory, allowing them to be restored later.

        Args:
            paged_ssd_cache_manager: PagedSSDCacheManager instance
        """
        self._paged_ssd_cache_manager = paged_ssd_cache_manager
        logger.info("paged SSD cache manager connected to PagedCacheManager")

    # =========================================================================
    # Dynamic Block Pool Growth (Elastic KV Cache)
    # =========================================================================

    def _grow_blocks(self, additional_blocks: int) -> int:
        """
        Dynamically expand the block pool.

        This method is called when more blocks are needed but the current
        pool is exhausted. It creates new blocks up to max_blocks limit.

        Args:
            additional_blocks: Number of blocks to add.

        Returns:
            Number of blocks actually created.
        """
        with self._lock:
            available = self.max_blocks - self._current_allocated_count
            to_create = min(additional_blocks, available)

            if to_create <= 0:
                return 0

            start_id = self._current_allocated_count
            new_blocks = [
                CacheBlock(block_id=i)
                for i in range(start_id, start_id + to_create)
            ]

            self.blocks.extend(new_blocks)
            self.free_block_queue.append_n(new_blocks)
            self._current_allocated_count += to_create

            self.stats.total_blocks = self._current_allocated_count
            self.stats.free_blocks = self.free_block_queue.num_free_blocks

            logger.info(
                f"Block pool grown: +{to_create}, "
                f"total={self._current_allocated_count}/{self.max_blocks}"
            )
            return to_create

    # =========================================================================
    # Block Allocation (vLLM style)
    # =========================================================================

    def allocate_block(self) -> Optional[CacheBlock]:
        """
        Allocate a new cache block.

        Returns:
            CacheBlock if available, None if out of memory.
        """
        with self._lock:
            if self.free_block_queue.num_free_blocks == 0:
                # Try to grow the block pool dynamically
                grown = self._grow_blocks(min(256, self.max_blocks - self._current_allocated_count))
                if grown == 0:
                    logger.warning("Out of cache blocks (max reached)")
                    return None

            block = self.free_block_queue.popleft()

            # Evict from hash cache if needed
            if self.enable_caching:
                self._maybe_evict_cached_block(block)

            block.ref_count = 1
            block.touch()
            self.allocated_blocks[block.block_id] = block

            self.stats.allocated_blocks += 1
            self.stats.free_blocks -= 1

            return block

    def get_new_blocks(self, num_blocks: int) -> List[CacheBlock]:
        """
        Allocate multiple blocks at once (vLLM style).

        Args:
            num_blocks: Number of blocks to allocate

        Returns:
            List of allocated blocks

        Raises:
            ValueError: If not enough free blocks even after dynamic growth
        """
        with self._lock:
            if num_blocks > self.free_block_queue.num_free_blocks:
                # Try to grow the block pool dynamically
                needed = num_blocks - self.free_block_queue.num_free_blocks
                self._grow_blocks(needed + 128)  # Extra buffer for future allocations

            if num_blocks > self.free_block_queue.num_free_blocks:
                raise ValueError(
                    f"Cannot allocate {num_blocks} blocks, "
                    f"only {self.free_block_queue.num_free_blocks} available "
                    f"(max={self.max_blocks})"
                )

            blocks = self.free_block_queue.popleft_n(num_blocks)

            for block in blocks:
                if self.enable_caching:
                    self._maybe_evict_cached_block(block)

                block.ref_count = 1
                block.touch()
                self.allocated_blocks[block.block_id] = block

            self.stats.allocated_blocks += num_blocks
            self.stats.free_blocks -= num_blocks

            return blocks

    def _notify_hash_dropped(self, block_hash: BlockHash | None) -> None:
        """Fire on_block_hash_dropped once a hash maps to no block at all.

        Called from every path that can kill a (hash -> block) association —
        the internal free/evict paths hold self._lock, while the prefix
        cache's own map pops call it lock-free (dict reads and the callback's
        dict.pop are GIL-atomic). The get_block re-check keeps hybrid models
        correct: the same hash can map to several blocks (one per KV cache
        group), and the association only dies with the last one.
        """
        if block_hash is None or self.on_block_hash_dropped is None:
            return
        if self.cached_block_hash_to_block.get_block(block_hash) is None:
            self.on_block_hash_dropped(block_hash)

    def _maybe_evict_cached_block(self, block: CacheBlock) -> bool:
        """
        Evict a block from the hash cache if present.

        In paged SSD-only mode, block data is always on paged SSD, so this just
        removes the block from the cache index.

        Args:
            block: Block to evict

        Returns:
            True if block was evicted from cache
        """
        if block.block_hash is None:
            return False

        block_hash = block.block_hash
        evicted = self.cached_block_hash_to_block.pop(block_hash, block.block_id)
        self._notify_hash_dropped(block_hash)

        if evicted:
            block.reset_hash()
            self.stats.evictions += 1
            return True

        return False

    def free_block(self, block_id: int) -> bool:
        """
        Free a cache block (decrements ref_count, frees if 0).

        Returns:
            True if block was freed, False if still referenced.
        """
        with self._lock:
            if block_id not in self.allocated_blocks:
                logger.warning(f"Attempted to free unknown block: {block_id}")
                return False

            block = self.allocated_blocks[block_id]
            if block.is_null:
                return False  # Never free null block

            block.ref_count -= 1

            if block.ref_count <= 0:
                # Remove from hash cache
                if block.block_hash is not None:
                    self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
                    self._notify_hash_dropped(block.block_hash)

                # Remove from allocated
                del self.allocated_blocks[block_id]

                # Add to free queue (back = MRU)
                self.free_block_queue.append(block)

                self.stats.allocated_blocks -= 1
                self.stats.free_blocks += 1
                self.stats.total_tokens_cached -= block.token_count

                return True

            return False

    def free_blocks(self, blocks: Iterable[CacheBlock]) -> None:
        """
        Free multiple blocks (vLLM style).

        Blocks with ref_count=0 are added to the free queue.

        Args:
            blocks: Blocks to free (in eviction order)
        """
        with self._lock:
            blocks_list = list(blocks)
            to_free = []

            for block in blocks_list:
                if block.is_null:
                    continue

                block.ref_count -= 1

                if block.ref_count <= 0:
                    # Remove from hash cache
                    if block.block_hash is not None:
                        self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
                        self._notify_hash_dropped(block.block_hash)

                    del self.allocated_blocks[block.block_id]
                    to_free.append(block)
                    self.stats.allocated_blocks -= 1
                    self.stats.free_blocks += 1
                    self.stats.total_tokens_cached -= block.token_count

            # Add to free queue (back = MRU, evicted last)
            self.free_block_queue.append_n(to_free)

    def touch(self, blocks: Iterable[CacheBlock]) -> None:
        """
        Touch blocks to prevent eviction (cache hit, vLLM style).

        Increments ref_count and removes from free queue if needed.

        Args:
            blocks: Blocks to touch
        """
        with self._lock:
            for block in blocks:
                if block.ref_count == 0 and not block.is_null:
                    # Block is in free queue, remove it
                    try:
                        self.free_block_queue.remove(block)
                        self.stats.free_blocks -= 1
                        self.stats.allocated_blocks += 1
                        self.allocated_blocks[block.block_id] = block
                    except RuntimeError:
                        pass  # Block not in queue

                block.ref_count += 1
                block.touch()

    # =========================================================================
    # Reference Counting
    # =========================================================================

    def increment_ref(self, block_id: int) -> bool:
        """Increment reference count for a block."""
        with self._lock:
            if block_id not in self.allocated_blocks:
                return False

            block = self.allocated_blocks[block_id]
            block.ref_count += 1
            block.touch()

            if block.ref_count == 2:
                self.stats.shared_blocks += 1

            return True

    def acquire_cached_block(
        self, block_id: int, expected_hash: BlockHash
    ) -> CacheBlock | None:
        """Atomically take a reference on a block iff it still holds a hash.

        Prefix-index entries can outlive the blocks they point at (a block
        may have been freed and reused for other content). Checking the hash
        and taking the reference under the same lock keeps a concurrent
        eviction from slipping in between the two steps.

        Args:
            block_id: Physical block index to acquire.
            expected_hash: Chain hash the block must still hold.

        Returns:
            The block on success, None when it is gone or reassigned.
        """
        with self._lock:
            block = self.allocated_blocks.get(block_id)
            if (
                block is None
                or block.is_null
                or block.block_hash != expected_hash
            ):
                return None
            block.ref_count += 1
            block.touch()
            if block.ref_count == 2:
                self.stats.shared_blocks += 1
            return block

    def decrement_ref(self, block_id: int) -> bool:
        """Decrement reference count (alias for free_block)."""
        return self.free_block(block_id)

    def release_for_eviction(self, block_ids: List[int]) -> int:
        """
        Release blocks for eviction without removing from allocated_blocks.

        Decrements ref_count so blocks become evictable, but keeps them in
        allocated_blocks so they can still be found by cache lookups.

        Args:
            block_ids: List of block IDs to release.

        Returns:
            Number of blocks released.
        """
        released = 0
        with self._lock:
            for block_id in block_ids:
                block = self.allocated_blocks.get(block_id)
                if block is None or block.is_null:
                    continue

                if block.ref_count > 0:
                    block.ref_count -= 1
                    released += 1

                    if block.ref_count == 1:
                        self.stats.shared_blocks = max(0, self.stats.shared_blocks - 1)

        return released

    # =========================================================================
    # Prefix Caching (vLLM chain-hash style)
    # =========================================================================

    def get_cached_block(self, block_hash: BlockHash) -> Optional[CacheBlock]:
        """
        Get a cached block by its hash (vLLM style).

        Args:
            block_hash: Content hash of the block

        Returns:
            Cached block if found, None otherwise
        """
        if not self.enable_caching:
            return None

        with self._lock:
            block = self.cached_block_hash_to_block.get_block(block_hash)
            if block:
                self.stats.hits += 1
            else:
                self.stats.misses += 1
            return block

    def cache_full_blocks(
        self,
        blocks: List[CacheBlock],
        token_ids: List[int],
        num_cached_blocks: int,
        num_full_blocks: int,
        extra_keys: Optional[Tuple[Any, ...]] = None,
    ) -> None:
        """
        Cache full blocks for prefix caching (vLLM style).

        Computes chain hashes and adds blocks to the cache.

        Args:
            blocks: All blocks for the request
            token_ids: All token IDs for the request
            num_cached_blocks: Number of blocks already cached
            num_full_blocks: Number of full blocks to cache
            extra_keys: Additional keys for hash (e.g., VLM image hash)
        """
        if not self.enable_caching:
            return

        if num_cached_blocks >= num_full_blocks:
            return

        with self._lock:
            # Get parent hash from last cached block
            parent_hash = None
            if num_cached_blocks > 0:
                parent_hash = blocks[num_cached_blocks - 1].block_hash

            for i in range(num_cached_blocks, num_full_blocks):
                block = blocks[i]
                if block.block_hash is not None:
                    parent_hash = block.block_hash
                    continue  # Already cached

                # Get tokens for this block
                start = i * self.block_size
                end = start + self.block_size
                block_tokens = token_ids[start:end]

                # Compute chain hash
                block_hash = compute_block_hash(
                    parent_hash, block_tokens,
                    extra_keys=extra_keys, model_name=self.model_name,
                )
                block.block_hash = block_hash
                block.token_count = len(block_tokens)

                # Add to cache
                self.cached_block_hash_to_block.insert(block_hash, block)

                parent_hash = block_hash

    def get_computed_blocks(
        self,
        token_ids: List[int],
        extra_keys: Optional[Tuple[Any, ...]] = None,
        extra_key_token_start: Optional[int] = None,
        extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
    ) -> Tuple[List[CacheBlock], int]:
        """
        Find cached blocks for a token prefix (vLLM style).

        Args:
            token_ids: Token IDs to look up
            extra_keys: Additional keys for hash (e.g., VLM image hash)

        Returns:
            Tuple of (cached_blocks, num_cached_tokens)
        """
        if not self.enable_caching:
            return [], 0

        with self._lock:
            cached_blocks = []
            parent_hash = None
            num_cached_tokens = 0

            num_full_blocks = len(token_ids) // self.block_size

            for i in range(num_full_blocks):
                start = i * self.block_size
                end = start + self.block_size
                block_tokens = token_ids[start:end]
                block_extra_keys = resolve_block_extra_keys(
                    end,
                    extra_keys=extra_keys,
                    extra_key_token_start=extra_key_token_start,
                    extra_key_ranges=extra_key_ranges,
                )

                # Compute expected hash
                block_hash = compute_block_hash(
                    parent_hash, block_tokens,
                    extra_keys=block_extra_keys, model_name=self.model_name,
                )

                # Look up in cache
                cached_block = self.cached_block_hash_to_block.get_block(block_hash)

                # Lazy restore: if not in memory but exists on SSD, register it
                if cached_block is None and self._paged_ssd_cache_manager is not None:
                    if self._paged_ssd_cache_manager.has_block(block_hash):
                        # Use standard allocation path so we handle an empty
                        # free queue gracefully (grow/evict) and keep stats in sync.
                        block = self.allocate_block()
                        if block is not None:
                            block.block_hash = block_hash
                            block.token_count = self.block_size
                            # Cold-registered blocks are metadata-only until a
                            # request claims them via increment_ref().
                            block.ref_count = 0
                            self.cached_block_hash_to_block.insert(
                                block_hash, block
                            )
                            cached_block = block

                if cached_block is None:
                    self.stats.misses += 1
                    break  # Cache miss, stop here

                cached_blocks.append(cached_block)
                parent_hash = block_hash
                num_cached_tokens += self.block_size
                self.stats.hits += 1

            return cached_blocks, num_cached_tokens

    # =========================================================================
    # Legacy hash methods (for backwards compatibility)
    # =========================================================================

    def find_cached_block(
        self,
        tokens: List[int],
        parent_hash: Optional[BlockHash] = None,
        extra_keys: Optional[Tuple[Any, ...]] = None,
    ) -> Optional[CacheBlock]:
        """
        Find a cached block matching the given tokens using chain hash.

        Args:
            tokens: Token IDs to look up
            parent_hash: Hash of the parent block (for chain), or None for first block
            extra_keys: Additional keys for hash (e.g., VLM image hash)

        Returns:
            Cached block if found, None otherwise
        """
        if not self.enable_caching:
            return None

        with self._lock:
            block_hash = compute_block_hash(
                parent_hash, tokens, extra_keys=extra_keys,
                model_name=self.model_name,
            )
            block = self.cached_block_hash_to_block.get_block(block_hash)
            if block:
                block.touch()
                self.stats.hits += 1
                return block

            self.stats.misses += 1
            return None

    def register_block_hash(
        self,
        block: CacheBlock,
        tokens: List[int],
        parent_hash: Optional[BlockHash] = None,
        extra_keys: Optional[Tuple[Any, ...]] = None,
    ) -> None:
        """
        Register a block's hash for deduplication using chain hash.

        Args:
            block: Block to register
            tokens: Token IDs in this block
            parent_hash: Hash of the parent block (for chain), or None for first block
            extra_keys: Additional keys for hash (e.g., VLM image hash)
        """
        if not self.enable_caching:
            return

        with self._lock:
            block_hash = compute_block_hash(
                parent_hash, tokens, extra_keys=extra_keys,
                model_name=self.model_name,
            )
            block.block_hash = block_hash
            self.cached_block_hash_to_block.insert(block_hash, block)

    # =========================================================================
    # Block Table Management
    # =========================================================================

    def create_block_table(self, request_id: str) -> BlockTable:
        """Create a new block table for a request."""
        with self._lock:
            table = BlockTable(request_id=request_id)
            self.request_tables[request_id] = table
            return table

    def get_block_table(self, request_id: str) -> Optional[BlockTable]:
        """Get block table for a request."""
        with self._lock:
            return self.request_tables.get(request_id)

    def get_or_create_block_table(self, request_id: str) -> BlockTable:
        """Get or create block table for a request."""
        with self._lock:
            if request_id not in self.request_tables:
                self.request_tables[request_id] = BlockTable(request_id=request_id)
            return self.request_tables[request_id]

    def delete_block_table(self, request_id: str) -> None:
        """Delete block table and free associated blocks."""
        with self._lock:
            table = self.request_tables.pop(request_id, None)
            if table:
                for block_id in table.block_ids:
                    self.free_block(block_id)

    def add_block_to_table(
        self,
        table: BlockTable,
        block: CacheBlock,
        tokens_in_block: int,
    ) -> None:
        """Add a block to a block table."""
        with self._lock:
            table.block_ids.append(block.block_id)
            block.token_count = tokens_in_block
            table.num_tokens += tokens_in_block
            self.stats.total_tokens_cached += tokens_in_block

    # =========================================================================
    # Prefix Sharing & COW
    # =========================================================================

    def find_shared_prefix(
        self,
        tokens: List[int],
        extra_keys: Optional[Tuple[Any, ...]] = None,
        extra_key_token_start: Optional[int] = None,
        extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
    ) -> Tuple[List[int], List[int]]:
        """
        Find shared prefix blocks for a token sequence.

        Uses get_computed_blocks for consistent chain-hash lookup.
        """
        cached_blocks, num_cached_tokens = self.get_computed_blocks(
            tokens,
            extra_keys=extra_keys,
            extra_key_token_start=extra_key_token_start,
            extra_key_ranges=extra_key_ranges,
        )

        shared_block_ids = [b.block_id for b in cached_blocks]
        remaining_tokens = tokens[num_cached_tokens:]

        return shared_block_ids, remaining_tokens

    def fork_block_table(
        self,
        source_table: BlockTable,
        new_request_id: str,
    ) -> BlockTable:
        """
        Fork a block table for a new request (COW).
        """
        with self._lock:
            new_table = source_table.copy(new_request_id)

            for block_id in new_table.block_ids:
                self.increment_ref(block_id)

            self.request_tables[new_request_id] = new_table

            logger.debug(
                f"Forked block table: {source_table.request_id} -> {new_request_id}, "
                f"blocks={len(new_table.block_ids)}"
            )

            return new_table

    def get_blocks_for_generation(
        self,
        table: BlockTable,
    ) -> Tuple[List[CacheBlock], bool]:
        """
        Get blocks for generation, applying COW if needed.
        """
        with self._lock:
            blocks = []
            was_copied = False

            for i, block_id in enumerate(table.block_ids):
                block = self.allocated_blocks.get(block_id)
                if not block:
                    continue

                if block.is_shared():
                    new_block = self._cow_copy_block(block)
                    if new_block:
                        table.block_ids[i] = new_block.block_id
                        blocks.append(new_block)
                        was_copied = True
                        self.stats.cow_copies += 1
                    else:
                        blocks.append(block)
                else:
                    blocks.append(block)

                block.touch()

            return blocks, was_copied

    def _cow_copy_block(self, source_block: CacheBlock) -> Optional[CacheBlock]:
        """
        Create a copy of a block for COW.

        In paged SSD-only mode, we don't copy data - we just allocate a new block
        with the same metadata. The actual KV data will be loaded from paged SSD
        when needed.
        """
        new_block = self.allocate_block()
        if not new_block:
            return None

        new_block.token_count = source_block.token_count
        new_block.block_hash = source_block.block_hash

        source_block.ref_count -= 1
        if source_block.ref_count == 1:
            self.stats.shared_blocks -= 1

        logger.debug(
            f"COW copy: block {source_block.block_id} -> {new_block.block_id}"
        )

        return new_block

    # =========================================================================
    # Legacy allocation methods (for backwards compatibility)
    # =========================================================================

    def allocate_blocks_for_tokens(self, num_tokens: int) -> List[CacheBlock]:
        """Allocate enough blocks to hold num_tokens."""
        num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
        return self.get_new_blocks(num_blocks_needed)

    # =========================================================================
    # Eviction
    # =========================================================================

    def evict_lru_blocks(self, num_blocks: int) -> int:
        """
        Evict least recently used blocks.

        With the doubly linked list, LRU blocks are already at the front
        of the free queue. We just need to pop from front.
        """
        with self._lock:
            evicted = 0

            # Get evictable blocks from free queue (they're already LRU ordered)
            for _ in range(min(num_blocks, self.free_block_queue.num_free_blocks)):
                try:
                    block = self.free_block_queue.popleft()
                    self._maybe_evict_cached_block(block)
                    # Put back at end (now available for allocation)
                    self.free_block_queue.append(block)
                    evicted += 1
                except ValueError:
                    break

            if evicted > 0:
                logger.info(f"Evicted {evicted} LRU blocks from cache")

            return evicted

    def handle_memory_pressure(self, requested_blocks: int) -> bool:
        """Handle memory pressure by evicting blocks."""
        with self._lock:
            if self.free_block_queue.num_free_blocks >= requested_blocks:
                return True

            needed = requested_blocks - self.free_block_queue.num_free_blocks
            self.evict_lru_blocks(needed)

            return self.free_block_queue.num_free_blocks >= requested_blocks

    # =========================================================================
    # Statistics and Properties
    # =========================================================================

    @property
    def free_blocks(self) -> int:
        """Number of free blocks available."""
        return self.free_block_queue.num_free_blocks

    @property
    def usage(self) -> float:
        """Cache usage ratio (0.0 to 1.0)."""
        total = self.max_blocks - 1  # Exclude null block
        if total == 0:
            return 0.0
        return 1.0 - (self.free_blocks / total)

    def get_stats(self) -> PagedCacheStats:
        """Get current cache statistics."""
        with self._lock:
            self.stats.shared_blocks = sum(
                1 for b in self.allocated_blocks.values() if b.ref_count > 1
            )
            self.stats.free_blocks = self.free_block_queue.num_free_blocks
            return self.stats

    def get_memory_usage(self) -> Dict[str, Any]:
        """Get memory usage information."""
        with self._lock:
            stats = self.get_stats()
            return {
                "block_size": self.block_size,
                "max_blocks": self.max_blocks,
                "allocated_blocks": stats.allocated_blocks,
                "free_blocks": stats.free_blocks,
                "shared_blocks": stats.shared_blocks,
                "total_tokens_cached": stats.total_tokens_cached,
                "utilization": stats.allocated_blocks / self.max_blocks,
                "cache_hit_rate": (
                    stats.hits / (stats.hits + stats.misses)
                    if (stats.hits + stats.misses) > 0 else 0
                ),
            }

    def reset_stats(self) -> None:
        """Reset statistics counters."""
        with self._lock:
            self.stats.hits = 0
            self.stats.misses = 0
            self.stats.cow_copies = 0
            self.stats.evictions = 0

    def reset_prefix_cache(self) -> bool:
        """Reset the prefix cache."""
        with self._lock:
            num_used = self.max_blocks - self.free_block_queue.num_free_blocks
            if num_used > 1:  # null_block is always "used"
                logger.warning(f"Cannot reset cache: {num_used - 1} blocks in use")
                return False

            self.cached_block_hash_to_block.clear()
            if self.on_hash_map_cleared is not None:
                self.on_hash_map_cleared()

            for block in self.blocks:
                block.reset_hash()

            self.stats.evictions = 0
            self.stats.hits = 0
            self.stats.misses = 0

            logger.info("Prefix cache reset successfully")
            return True

    def clear(self) -> int:
        """
        Clear all cached data and reset to initial block count.

        Returns:
            Number of entries cleared.
        """
        with self._lock:
            # Count entries before clearing
            cleared_count = len(self.allocated_blocks) - 1  # Exclude null block

            # Reset to initial blocks (memory optimization)
            initial_count = min(self.initial_blocks, self.max_blocks)
            self._current_allocated_count = initial_count

            # Recreate blocks and queue with only initial blocks
            self.blocks = [
                CacheBlock(block_id=i) for i in range(initial_count)
            ]
            self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)

            self.cached_block_hash_to_block.clear()
            if self.on_hash_map_cleared is not None:
                self.on_hash_map_cleared()
            self.request_tables.clear()
            self.allocated_blocks.clear()

            # Reserve null block
            self.null_block = self.free_block_queue.popleft()
            self.null_block.is_null = True
            self.null_block.ref_count = 1
            self.allocated_blocks[self.null_block.block_id] = self.null_block

            self.stats = PagedCacheStats(
                total_blocks=initial_count,
                allocated_blocks=1,
                free_blocks=initial_count - 1,
            )

            logger.info(
                f"PagedCacheManager cleared (reset to {initial_count} initial blocks)"
            )

            return max(0, cleared_count)

    # =========================================================================
    # SSD Cache Support
    # =========================================================================

    def get_evictable_blocks(self, count: int) -> List[CacheBlock]:
        """
        Get LRU blocks that can be evicted (metadata cleared).

        In paged SSD-only mode, blocks don't store data in GPU memory,
        so this returns blocks that can be freed from the metadata index.

        Args:
            count: Maximum number of blocks to return.

        Returns:
            List of evictable blocks in LRU order.
        """
        with self._lock:
            candidates = []

            # Iterate through free queue (LRU order)
            current = self.free_block_queue.fake_head.next_free_block
            while (
                current is not None
                and current != self.free_block_queue.fake_tail
                and len(candidates) < count
            ):
                # Block must not be null and have ref_count == 0
                if not current.is_null and current.ref_count == 0:
                    candidates.append(current)
                current = current.next_free_block

            # Also check allocated blocks with ref_count == 0 (not in free queue yet)
            if len(candidates) < count:
                # Sort by last_access (LRU)
                remaining = []
                for block in self.allocated_blocks.values():
                    if (
                        not block.is_null
                        and block.ref_count == 0
                        and block not in candidates
                    ):
                        remaining.append(block)

                remaining.sort(key=lambda b: b.last_access)
                candidates.extend(remaining[: count - len(candidates)])

            return candidates[:count]

    def mark_block_cold(self, block_id: int) -> bool:
        """
        Mark a block as evicted (metadata preserved).

        In paged SSD-only mode, this is a no-op since block data is always on paged SSD.
        Kept for API compatibility.

        Args:
            block_id: Block ID to mark.

        Returns:
            True if successful, False if block not found or has data users.
        """
        with self._lock:
            block = self.blocks[block_id] if block_id < len(self.blocks) else None
            if block is None:
                logger.warning(f"Block {block_id} not found")
                return False

            if block.ref_count > 0:
                logger.warning(
                    f"Cannot mark block {block_id}: ref_count={block.ref_count}"
                )
                return False

            if block.is_null:
                logger.warning(f"Cannot mark null block")
                return False

            # In paged SSD-only mode, data is already on paged SSD
            self.stats.evictions += 1

            logger.debug(
                f"Marked block {block_id} "
                f"(hash={block.block_hash.hex()[:16] if block.block_hash else 'None'}...)"
            )
            return True

    def evict_block_permanently(self, block_id: int) -> bool:
        """
        Evict a block permanently (removes from metadata index).

        This method:
        - Removes from hash cache (block won't be found in cache lookups)
        - Returns block to free queue (can be reallocated)
        - Removes from allocated_blocks

        Note: In paged SSD-only mode, the data remains on paged SSD and may be deleted
        by PagedSSDCacheManager's LRU eviction if needed.

        Args:
            block_id: Block ID to evict.

        Returns:
            True if successful, False if block not found or in use.
        """
        with self._lock:
            block = self.blocks[block_id] if block_id < len(self.blocks) else None
            if block is None:
                logger.warning(f"Block {block_id} not found for permanent eviction")
                return False

            if block.ref_count > 0:
                logger.warning(
                    f"Cannot permanently evict block {block_id}: ref_count={block.ref_count}"
                )
                return False

            if block.is_null:
                logger.warning(f"Cannot evict null block")
                return False

            # Remove from hash cache
            if block.block_hash is not None:
                self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
                self._notify_hash_dropped(block.block_hash)

            # Clear metadata
            block.reset_hash()
            block.token_count = 0

            # Remove from allocated_blocks and add to free queue
            if block_id in self.allocated_blocks:
                del self.allocated_blocks[block_id]
                self.stats.allocated_blocks -= 1

            self.free_block_queue.append(block)
            self.stats.free_blocks += 1
            self.stats.evictions += 1

            logger.debug(f"Permanently evicted block {block_id}")
            return True

    def restore_block(
        self,
        block_id: int,
        cache_data: List[Tuple[Any, Any]],
    ) -> bool:
        """
        Restore block data from cold storage.

        In paged SSD-only mode, this is a no-op since data is always loaded
        directly from paged SSD when needed. Kept for API compatibility.

        Args:
            block_id: Block ID to restore.
            cache_data: KV cache data (unused in paged SSD-only mode).

        Returns:
            True if successful, False if block not found.
        """
        with self._lock:
            block = self.blocks[block_id] if block_id < len(self.blocks) else None
            if block is None:
                logger.warning(f"Block {block_id} not found for restoration")
                return False

            block.touch()

            logger.debug(
                f"Block {block_id} touched "
                f"(hash={block.block_hash.hex()[:16] if block.block_hash else 'None'}...)"
            )
            return True

    def get_cold_blocks(self) -> List[CacheBlock]:
        """
        Get all blocks that have data on paged SSD.

        In paged SSD-only mode, returns all blocks with block_hash set
        (i.e., blocks that have data stored on paged SSD).

        Returns:
            List of blocks with paged SSD data.
        """
        with self._lock:
            return [b for b in self.blocks if b.block_hash is not None and not b.is_null]

    @property
    def cold_block_count(self) -> int:
        """Number of blocks with data on paged SSD."""
        with self._lock:
            return sum(1 for b in self.blocks if b.block_hash is not None and not b.is_null)

    def get_ref_count_distribution(self) -> Dict[int, int]:
        """
        Get distribution of blocks by ref_count.

        Returns:
            Dict mapping ref_count -> number of blocks with that count.
            Only includes ref_counts that have at least one block.
        """
        with self._lock:
            distribution: Dict[int, int] = {}
            for block in self.allocated_blocks.values():
                rc = block.ref_count
                distribution[rc] = distribution.get(rc, 0) + 1
            return distribution

    def get_ref_count_summary(self) -> str:
        """
        Get a compact string summary of ref_count distribution.

        Returns:
            String like "rc0=5(ssd=3),rc1=100" showing counts per ref_count.
            In paged SSD-only mode, all blocks have data on paged SSD.
        """
        dist = self.get_ref_count_distribution()
        if not dist:
            return "rc=none"

        # Count blocks with paged SSD data (blocks with block_hash)
        ssd_count = 0
        with self._lock:
            for block in self.allocated_blocks.values():
                if block.ref_count == 0 and block.block_hash is not None:
                    ssd_count += 1

        parts = []
        for k, v in sorted(dist.items()):
            if k == 0:
                # Show paged SSD count for rc0
                parts.append(f"rc0={v}(ssd={ssd_count})")
            else:
                parts.append(f"rc{k}={v}")
        return ",".join(parts)

    # =========================================================================
    # CacheManager ABC Interface Implementation
    # =========================================================================

    def fetch(self, key: Any) -> Tuple[Optional[Any], bool]:
        """
        Fetch a cached block by its hash.

        Args:
            key: BlockHash (bytes) to look up.

        Returns:
            Tuple of (CacheBlock, True) if found, (None, False) otherwise.
        """
        if not isinstance(key, bytes):
            return None, False

        block = self.get_cached_block(BlockHash(key))
        if block is not None:
            return block, True
        return None, False

    def store(self, key: Any, value: Any) -> bool:
        """
        Store a block in the cache.

        For PagedCacheManager, use allocate_block() and register_block_hash()
        for the full workflow. This method provides a simplified interface.

        Args:
            key: BlockHash (bytes) for the block.
            value: CacheBlock to store.

        Returns:
            True if stored successfully.
        """
        if not isinstance(key, bytes) or not isinstance(value, CacheBlock):
            return False

        with self._lock:
            block_hash = BlockHash(key)
            value.block_hash = block_hash
            self.cached_block_hash_to_block.insert(block_hash, value)
            return True

    def evict(self, key: Any) -> bool:
        """
        Evict a specific block from the cache.

        Args:
            key: BlockHash (bytes) or block_id (int) to evict.

        Returns:
            True if evicted, False if not found.
        """
        if isinstance(key, bytes):
            # Evict by block hash
            block = self.cached_block_hash_to_block.get_block(BlockHash(key))
            if block is not None:
                return self.evict_block_permanently(block.block_id)
            return False
        elif isinstance(key, int):
            # Evict by block ID
            return self.evict_block_permanently(key)
        return False

    @property
    def size(self) -> int:
        """
        Get the current number of allocated blocks.

        Returns:
            Number of allocated blocks (excluding null block).
        """
        return max(0, len(self.allocated_blocks) - 1)  # Exclude null block

    @property
    def max_size(self) -> int:
        """
        Get the maximum number of blocks.

        Returns:
            Maximum number of blocks.
        """
        return self.max_blocks


# ===== FILE: omlx/cache/paged_ssd_cache.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Paged SSD Cache Manager for oMLX KV cache.

This module implements SSD-based storage for paged KV cache blocks,
enabling larger effective cache sizes than GPU memory allows.

Key features:
- Block-level safetensors serialization (compatible with mlx-lm)
- Hash-based subdirectory structure for scalability
- LRU-based paged SSD cache size management
- Startup scan to reuse existing cache files

Reference: mlx-lm/mlx_lm/models/cache.py (save_prompt_cache, load_prompt_cache)
"""

from __future__ import annotations

import contextlib
import errno
import json
import logging
import os
import queue
import shutil
import struct
import threading
import time
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from omlx.utils.formatting import format_bytes

from .interface import CacheManager
from .pooling_delta import (
    POOLING_CACHE_DELTA_CLASS,
    POOLING_CACHE_DELTA_FORMAT_VERSION,
)
from .stats import PagedSSDCacheStats

logger = logging.getLogger(__name__)

# Check for MLX
try:
    import mlx.core as mx
    from mlx.utils import tree_flatten, tree_unflatten

    HAS_MLX = True
except ImportError:
    HAS_MLX = False
    mx = None


# --- Async I/O constants ---
# Fraction of host RAM the pending-write queue targets at saturation.
# The queue holds raw-byte copies of KV blocks that the background
# writer hasn't drained yet (see ``_extract_tensor_bytes`` in
# ``save_block``). The hard fraction below bounds the soft floor so
# large-block workloads cannot silently reserve an unsafe amount of RAM.
_PENDING_WRITES_TARGET_RAM_FRACTION = 0.10
_PENDING_WRITES_HARD_RAM_FRACTION = 0.30
_PENDING_WRITES_SOFT_FLOOR = 32
_PENDING_WRITES_CEILING = 256
_PENDING_WRITE_PUT_TIMEOUT_SECONDS = 1.0

# Conservative defaults for the per-block cost estimator. The actual
# bytes-per-block depends on the model (num_layers × num_kv_heads ×
# head_dim × dtype_size × block_size_tokens × 2). At construction time
# the PagedSSDCacheManager doesn't always know these — see __init__'s
# ``expected_kv_bytes_per_token`` parameter — so the module-level
# default targets a 35B-class bf16 model whose per-token KV is ≈200 KB
# spread across all layers. Smaller models will be over-conservative
# (fine), larger models or larger blocks should pass an explicit value.
_DEFAULT_BLOCK_SIZE_TOKENS = 256
_DEFAULT_KV_BYTES_PER_TOKEN = 200_000


def _compute_max_pending_writes(
    block_size_tokens: int = _DEFAULT_BLOCK_SIZE_TOKENS,
    kv_bytes_per_token: int = _DEFAULT_KV_BYTES_PER_TOKEN,
    target_fraction: float = _PENDING_WRITES_TARGET_RAM_FRACTION,
    hard_fraction: float = _PENDING_WRITES_HARD_RAM_FRACTION,
) -> int:
    """Compute max pending writes queue depth.

    Scales by *block bytes* so the target pending pool stays near
    ``target_fraction`` of host RAM regardless of how big each block is,
    while ``hard_fraction`` bounds the soft floor:

        worst_case_bytes = cap × block_size_tokens × kv_bytes_per_token
        cap = (total_ram × target_fraction) / (block_size × kv_bytes_per_token)

    Bounded by a soft floor, a byte hard cap, and a ceiling:
      - Soft floor at 32 so even small systems with large blocks retain
        burst headroom for a few in-flight writes — dropping to zero
        means every save serializes against the disk and the writer
        thread becomes a hard bottleneck on the inference loop.
      - Hard cap at 30% of host RAM so the soft floor cannot turn very
        large blocks into an unsafe memory reservation.
      - Ceiling at 256 so 512 GB+ systems don't pin gigabytes against
        a writer that's already keeping up at lower caps.

    Workload sizing: long-context coding workloads snapshot ~73 blocks
    per turn at 150 K tokens (block_size=2048), and at the default
    block size of 256 tokens that's ~586 blocks per snapshot. The
    queue is a *burst ceiling*, not steady state — a healthy writer
    drains it continuously and the cap only matters when the writer
    is fighting memory pressure or slow disk. Sustained saturation now
    falls back to inline writes, so the cap controls when save latency
    moves onto the request path.

    Defaults target a 35B-class bf16 model at the default
    ``paged_cache_block_size=256``; pass an explicit
    ``kv_bytes_per_token`` for larger models or quantized configs.
    """
    try:
        total_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
        block_bytes = max(1, block_size_tokens * kv_bytes_per_token)
        target = int(total_bytes * target_fraction / block_bytes)
        hard_cap = max(1, int(total_bytes * hard_fraction / block_bytes))
        soft_target = max(_PENDING_WRITES_SOFT_FLOOR, target)
        return max(1, min(_PENDING_WRITES_CEILING, soft_target, hard_cap))
    except (ValueError, OSError):
        return 64  # Safe default


# Module-level constant for static callers that don't have model
# config. The PagedSSDCacheManager recomputes per-instance from its
# ``__init__`` parameters so a non-default block size or model
# generation can plumb through.
_MAX_PENDING_WRITES = _compute_max_pending_writes()

# Cap on the number of LRU blocks ``_enforce_size_limit_for_new_block`` is
# allowed to unlink in one inline burst. Eviction normally returns ~1
# entry; the cap exists for the ENOSPC-recovery path where the disk-usage
# cache invalidates and the next ``_get_effective_max_size`` call can
# shrink sharply — ``evict_until_size`` would then return hundreds of
# entries at once and stall the inference thread on a syscall storm.
# Deferred-but-not-unlinked entries are reinserted into the index so
# subsequent saves drain the remainder; bounds per-call latency at the
# cost of taking multiple saves to fully reconverge.
_MAX_INLINE_UNLINKS_PER_SAVE = 32


# Cache format version. Bump when on-disk layout or RotatingKVCache meta_state
# semantics change in a way that older blocks become unsafe to load.
#
# Version "2": added with the mlx-lm 0.31.3 contract fix (issues #934 / #903).
# Version "1" / unset: pre-fix blocks. RotatingKVCache layers may have been
#   zero-padded to max_size, which after the fix would leak zero positions
#   into attention. Treat such blocks as a cache miss instead of migrating.
_CACHE_FORMAT_VERSION = "3"

# Versions whose blocks the current code can read. V3 polyfills V2 blocks
# whose layer data was stored as the legacy 2-tuple `(keys, values)` —
# they are upgraded to N-tuple markers on read so the rest of omlx core
# sees a uniform shape. New writes use V3 unless a PoolingCache block carries
# the V4 append-only delta representation.
_READABLE_CACHE_FORMAT_VERSIONS = frozenset(
    {"2", "3", POOLING_CACHE_DELTA_FORMAT_VERSION}
)


# Layer cache type names whose meta_state should be clamped on save so the
# rotating buffer's _idx never exceeds the actual buffer length. Restoring a
# cache where _idx > keys.shape[2] makes BatchRotatingKVCache.merge() either
# overshoot the RHS or (when omlx pads) leak zero positions into attention.
_ROTATING_CACHE_TYPES = (
    "RotatingKVCache",
    "BatchRotatingKVCache",
    "BufferedRotatingKVCache",
)

_STORAGE_CLASS_NAME_ALIASES = {
    # mlx-vlm MTP uses this transient target-cache wrapper for rollback
    # slack. Persist the canonical rotating-cache type so blocks remain
    # compatible with the model.make_cache() signature.
    "BufferedRotatingKVCache": "RotatingKVCache",
}


def _storage_layer_cache_types(
    layer_cache_types: list[str] | tuple[str, ...] | None,
) -> list[str] | None:
    """Return type names to persist in new cache block metadata."""
    if layer_cache_types is None:
        return None
    return [
        _STORAGE_CLASS_NAME_ALIASES.get(cache_type, cache_type)
        for cache_type in layer_cache_types
    ]


def _canonicalize_layer_cache_types(
    layer_cache_types: list[str] | tuple[str, ...] | None,
) -> list[str] | None:
    """Normalize wrapper class names for metadata compatibility checks.

    Wrapper classes that keep the same tensor representation compare equal.
    Types that change tensor representation (e.g., ``TurboQuantKVCache`` vs
    ``KVCache``) are NOT collapsed -- that mismatch is real and the block must
    be invalidated.
    """
    if layer_cache_types is None:
        return None
    wrapper_to_canonical = {
        "SizedArraysCache": "ArraysCache",
        "PrefillReadyRotatingKVCache": "RotatingKVCache",
        POOLING_CACHE_DELTA_CLASS: "PoolingCache",
        # Batch and single-request TurboQuant caches persist the same packed
        # per-request state (the save path records whichever class name it
        # extracted; the restore path rebuilds a TurboQuantKVCache from
        # either). Collapsing them keeps the predicted layout from
        # refresh_ssd_layer_signature — which always says
        # "TurboQuantKVCache" — from sweeping valid batch-form blocks.
        "BatchTurboQuantKVCache": "TurboQuantKVCache",
    }
    return [
        wrapper_to_canonical.get(cache_type, cache_type)
        for cache_type in layer_cache_types
    ]


def _cache_compat_signature(
    *,
    model_name: str = "",
    num_layers: int = 0,
    block_size: int = 0,
    layer_cache_types: list[str] | None = None,
    turboquant_kv_bits: float | None = None,
    cachelist_subtypes: dict[str, list[str]] | None = None,
) -> str:
    """Return a stable compatibility signature for a persisted cache block."""
    payload = {
        "model_name": model_name or "",
        "num_layers": int(num_layers or 0),
        "block_size": int(block_size or 0),
        "layer_cache_types": list(layer_cache_types or []),
    }
    # TurboQuant packed state width depends on the bit depth
    # (packed_width = ceil(head_dim * bits / 32)), so blocks written at
    # different bit depths are shape-incompatible (#2045). Only stamped
    # when TurboQuant is active so non-TurboQuant signatures stay
    # byte-identical to the previous format.
    if turboquant_kv_bits is not None:
        payload["turboquant_kv_bits"] = float(turboquant_kv_bits)
    # Mixed CacheList layers (a non-sliceable sub next to a KVCache, e.g.
    # inkling's CacheList(KVCache, ArraysCache(4))) additionally stamp
    # their sub composition: the flat "CacheList" type name cannot tell a
    # 2-slot ArraysCache block from a 4-slot one, and restoring the wrong
    # arity IndexErrors in the model. Only stamped when such layers exist
    # so other signatures stay byte-identical to the previous format.
    if cachelist_subtypes:
        payload["cachelist_subtypes"] = cachelist_subtypes
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def _signature_turboquant_bits(cache_signature: str) -> float | None:
    """Extract ``turboquant_kv_bits`` from a stored signature, or None."""
    if not cache_signature:
        return None
    try:
        payload = json.loads(cache_signature)
    except (TypeError, ValueError):
        return None
    if not isinstance(payload, dict):
        # Corrupted/foreign signature that parses as a JSON scalar or list.
        # Report "no recorded depth" instead of raising: an AttributeError
        # here would abort the whole stale-signature sweep.
        return None
    bits = payload.get("turboquant_kv_bits")
    if bits is None:
        return None
    try:
        return float(bits)
    except (TypeError, ValueError):
        return None


def _block_turboquant_bits(
    layer_cache_types: list[str] | None,
    layer_meta_states: list[tuple] | None,
) -> float | None:
    """Read the bit depth a block's own TurboQuant layers were packed at.

    TurboQuant meta_state is ``(offset, bits, seed, ...)`` — the same tuple
    the restore path reads back at reconstruction. Deriving the signature
    stamp from the block itself keeps it truthful even when the manager's
    expectation is stale or not yet learned.
    """
    if not layer_cache_types or not layer_meta_states:
        return None
    for i, cache_type in enumerate(layer_cache_types):
        if cache_type not in ("TurboQuantKVCache", "BatchTurboQuantKVCache"):
            continue
        if i >= len(layer_meta_states):
            continue
        meta_state = layer_meta_states[i]
        if isinstance(meta_state, (list, tuple)) and len(meta_state) >= 3:
            try:
                return float(meta_state[1])
            except (TypeError, ValueError):
                return None
    return None


_CACHELIST_NON_SLICEABLE_SUB_CLASSES = frozenset(
    {
        "ArraysCache",
        "SizedArraysCache",
        "PoolingCache",
        "BatchPoolingCache",
        "RotatingKVCache",
        "BatchRotatingKVCache",
        "PrefillReadyRotatingKVCache",
        "BufferedRotatingKVCache",
    }
)

_ARRAYS_SUB_CLASSES = frozenset({"ArraysCache", "SizedArraysCache"})
_POOLING_SUB_CLASSES = frozenset({"PoolingCache", "BatchPoolingCache"})


def _canonical_sub_name(name: Any) -> str:
    """Canonicalize one CacheList sub-cache class name for signatures."""
    canonical = _canonicalize_layer_cache_types([str(name or "")])
    return canonical[0] if canonical else ""


def _signature_cachelist_subtypes(cache_signature: str) -> dict | None:
    """Extract ``cachelist_subtypes`` from a stored signature, or None."""
    if not cache_signature:
        return None
    try:
        payload = json.loads(cache_signature)
    except (TypeError, ValueError):
        return None
    if not isinstance(payload, dict):
        return None
    subtypes = payload.get("cachelist_subtypes")
    return subtypes if isinstance(subtypes, dict) else None


def _block_cachelist_subtypes(
    cache_data: list[Any] | None,
    layer_cache_types: list[str] | None,
    layer_meta_states: list[tuple] | None,
) -> dict[str, list[str]] | None:
    """Describe mixed CacheList layers' sub composition from block payload.

    Only layers whose composition contains a non-sliceable sub class are
    stamped, so signatures of KVCache-only CacheList models (GLM /
    deepseek_v32) stay byte-identical to the previous format. ArraysCache
    descriptors carry the slot count (``"ArraysCache:4"``), while pooling
    descriptors carry the logical state arity (``"PoolingCache:5"``).
    Both values are part of the persisted layout identity.
    """
    if not cache_data or not layer_cache_types:
        return None
    subtypes: dict[str, list[str]] = {}
    for i, cache_type in enumerate(layer_cache_types):
        if cache_type != "CacheList" or i >= len(cache_data):
            continue
        layer_data = cache_data[i]
        if not (
            isinstance(layer_data, tuple)
            and len(layer_data) == 2
            and layer_data[0] == "__cache_list__"
            and isinstance(layer_data[1], (list, tuple))
        ):
            continue
        meta_names: list[Any] = []
        if (
            layer_meta_states
            and i < len(layer_meta_states)
            and isinstance(layer_meta_states[i], (list, tuple))
            and len(layer_meta_states[i]) >= 1
            and isinstance(layer_meta_states[i][0], (list, tuple))
        ):
            meta_names = list(layer_meta_states[i][0])
        descriptors: list[str] = []
        has_non_sliceable = False
        for j, sub_tensor in enumerate(layer_data[1]):
            name = _canonical_sub_name(meta_names[j] if j < len(meta_names) else "")
            element_count = None
            is_pooling_delta = False
            if (
                isinstance(sub_tensor, tuple)
                and len(sub_tensor) >= 3
                and sub_tensor[0] == "__nstate__"
            ):
                if not name:
                    name = _canonical_sub_name(sub_tensor[1])
                if isinstance(sub_tensor[2], (list, tuple)):
                    element_count = len(sub_tensor[2])
                    is_pooling_delta = sub_tensor[1] == POOLING_CACHE_DELTA_CLASS
            elif isinstance(sub_tensor, (list, tuple)):
                element_count = len(sub_tensor)
            if name in _CACHELIST_NON_SLICEABLE_SUB_CLASSES:
                has_non_sliceable = True
            if name in _ARRAYS_SUB_CLASSES and element_count is not None:
                descriptors.append(f"ArraysCache:{element_count}")
            elif name in _POOLING_SUB_CLASSES and element_count is not None:
                logical_count = element_count - 1 if is_pooling_delta else element_count
                descriptors.append(f"{name}:{logical_count}")
            else:
                descriptors.append(name or "?")
        if descriptors and has_non_sliceable:
            subtypes[str(i)] = descriptors
    return subtypes or None


def cachelist_subtypes_from_cache_list(
    cache_list: list[Any] | tuple[Any, ...] | None,
) -> dict[str, list[str]] | None:
    """Describe mixed CacheList layers' sub composition from live caches.

    The live-model counterpart of ``_block_cachelist_subtypes`` — produces
    the expectation the manager compares stored blocks against. Stamps the
    same layers (composition contains a non-sliceable sub) with the same
    descriptor format.
    """
    if not cache_list:
        return None
    subtypes: dict[str, list[str]] = {}
    for i, cache_obj in enumerate(cache_list):
        sub_caches = getattr(cache_obj, "caches", None)
        if type(cache_obj).__name__ != "CacheList" or not sub_caches:
            continue
        descriptors: list[str] = []
        has_non_sliceable = False
        for sub in sub_caches:
            name = _canonical_sub_name(type(sub).__name__)
            if name in _CACHELIST_NON_SLICEABLE_SUB_CLASSES:
                has_non_sliceable = True
            if name in _ARRAYS_SUB_CLASSES:
                slots = getattr(sub, "cache", None)
                slot_count = len(slots) if isinstance(slots, list) else 0
                descriptors.append(f"ArraysCache:{slot_count}")
            elif name in _POOLING_SUB_CLASSES:
                state = getattr(sub, "state", ())
                state_count = len(state) if isinstance(state, (list, tuple)) else 0
                descriptors.append(f"{name}:{state_count}")
            else:
                descriptors.append(name or "?")
        if descriptors and has_non_sliceable:
            subtypes[str(i)] = descriptors
    return subtypes or None


def _clamp_rotating_meta_states(
    cache_data: list[Any],
    layer_cache_types: list[str] | None,
    layer_meta_states: list[tuple] | None,
) -> list[tuple] | None:
    """Clamp ``_idx`` to ``keys.shape[2]`` for RotatingKVCache layers.

    RotatingKVCache.meta_state is ``(keep, max_size, offset, _idx)``. When
    we save a snapshot, ``_idx`` must reflect the actual buffer length so
    the restored cache lands in case 1 of ``_temporal_order``. Older code
    paths could leave ``_idx == max_size`` after zero-padding the buffer;
    by clamping at write time we ensure newer blocks are always safe to
    restore.
    """
    if not layer_meta_states or not layer_cache_types:
        return layer_meta_states

    clamped: list[tuple] = []
    for i, meta in enumerate(layer_meta_states):
        if (
            i < len(layer_cache_types)
            and layer_cache_types[i] in _ROTATING_CACHE_TYPES
            and meta
            and len(meta) >= 4
            and i < len(cache_data)
        ):
            layer_data = cache_data[i]
            seq_len: int | None = None
            if (
                isinstance(layer_data, tuple)
                and len(layer_data) == 2
                and not (
                    isinstance(layer_data[0], str) and layer_data[0].startswith("__")
                )
            ):
                keys = layer_data[0]
                if hasattr(keys, "shape") and len(keys.shape) >= 3:
                    seq_len = int(keys.shape[2])
            if seq_len is not None:
                try:
                    keep, max_size, offset, idx = meta[:4]
                    idx_int = int(idx)
                    if idx_int > seq_len:
                        clamped.append((keep, max_size, offset, str(seq_len)))
                        continue
                except (TypeError, ValueError):
                    pass
        clamped.append(meta)
    return clamped


def _has_zero_dim(tensor: Any) -> bool:
    """Check if a tensor has any zero-dimension axis (unsupported by safetensors)."""
    return hasattr(tensor, "shape") and any(d == 0 for d in tensor.shape)


def _encode_shape(shape) -> str:
    """Encode tensor shape as comma-separated string for safetensors metadata."""
    return ",".join(str(d) for d in shape)


def _decode_shape(shape_str: str) -> tuple:
    """Decode shape string back to tuple of ints."""
    return tuple(int(d) for d in shape_str.split(","))


def _store_nstate_elements_flat(
    arrays: dict[str, Any],
    cache_list_meta: dict[str, str],
    prefix: str,
    elements,
) -> None:
    """Write N elements as ``{prefix}_state_{k}`` keys with a
    ``{prefix}_state_count`` count marker. Zero-dim shapes are
    preserved via ``{prefix}_state_{k}_zero_dim``. Composite
    elements (a bare tuple/list or a nested ``__nstate__``
    marker) recurse under a ``{elem_key}`` sub-prefix and record
    a ``{elem_key}_nested`` marker; the flat ``{elem_key}``
    tensor is deliberately omitted so an older reader hits its
    ``Missing {elem_key} in arrays`` path and skips the block.

    Module-level on purpose — same recursive-closure cycle bug as
    ``_load_nstate_flat`` (see that docstring).
    """
    cache_list_meta[f"{prefix}_state_count"] = str(len(elements))
    for k, elem in enumerate(elements):
        elem_key = f"{prefix}_state_{k}"
        if elem is None:
            # None placeholder — store an empty marker tensor
            # and a sentinel zero_dim entry so the loader can
            # restore None instead of materializing zeros.
            arrays[elem_key] = mx.zeros((1,))
            cache_list_meta[f"{elem_key}_none"] = "1"
        elif _has_zero_dim(elem):
            arrays[elem_key] = mx.zeros((1,))
            cache_list_meta[f"{elem_key}_zero_dim"] = _encode_shape(elem.shape)
        elif (
            isinstance(elem, tuple)
            and len(elem) >= 2
            and isinstance(elem[0], str)
            and elem[0] == "__nstate__"
        ):
            # Nested ``('__nstate__', class_name, [sub...])``
            # marker — recurse, no flat tensor written.
            cache_list_meta[f"{elem_key}_nested"] = "nstate"
            sub_class = elem[1] if len(elem) >= 2 else None
            sub_elements = elem[2] if len(elem) >= 3 else []
            if sub_class:
                cache_list_meta[f"{elem_key}_state_class_name"] = sub_class
            _store_nstate_elements_flat(arrays, cache_list_meta, elem_key, sub_elements)
        elif isinstance(elem, (tuple, list)):
            # Bare tuple/list of sub-elements — recurse, no flat
            # tensor written.
            cache_list_meta[f"{elem_key}_nested"] = "tuple"
            _store_nstate_elements_flat(arrays, cache_list_meta, elem_key, list(elem))
        else:
            if not isinstance(elem, mx.array):
                raise TypeError(
                    f"unsupported non-array nstate element "
                    f"{elem_key}: {type(elem).__name__}"
                )
            arrays[elem_key] = elem


def _load_nstate_flat(
    arrays: dict[str, Any],
    file_metadata: dict[str, str],
    prefix: str,
    fallback_class: str | None,
) -> tuple | None:
    """Read either V3 ``state_count`` keys or V2 ``keys``/``values``
    polyfill at ``prefix``. Returns ``('__nstate__', class_name, elements)``
    on success or None on missing tensors.

    Module-level on purpose: the previous nested-closure version formed a
    self-referential cycle (recursive closure) that captured ``arrays`` —
    hundreds of MB of KV tensors — and only gen-2 gc could free it.
    """
    count_key = f"{prefix}_state_count"
    class_name = None
    if file_metadata:
        class_name = file_metadata.get(f"{prefix}_state_class_name")
    if class_name is None:
        class_name = fallback_class

    elements: list[Any] = []
    if file_metadata and count_key in file_metadata:
        # V3 path
        try:
            count = int(file_metadata[count_key])
        except (ValueError, TypeError):
            return None
        for k in range(count):
            elem_key = f"{prefix}_state_{k}"
            none_marker = f"{elem_key}_none"
            zd_marker = f"{elem_key}_zero_dim"
            nested_marker = f"{elem_key}_nested"
            if file_metadata and none_marker in file_metadata:
                elements.append(None)
                continue
            if file_metadata and nested_marker in file_metadata:
                # Composite element — recurse, then restore the same
                # shape it had on save (bare tuple vs __nstate__).
                sub = _load_nstate_flat(arrays, file_metadata, elem_key, None)
                if sub is None:
                    return None
                if file_metadata[nested_marker] == "tuple":
                    elements.append(tuple(sub[2]))
                elif file_metadata[nested_marker] == "nstate":
                    # Explicit marker on write — preserve the full
                    # ('__nstate__', class_name, elements) as-is;
                    # never unwrap (would drop the marker/class_name).
                    elements.append(sub)
                else:
                    # Corrupt/unknown nested marker — fail closed.
                    return None
                continue
            if elem_key not in arrays:
                logger.error(f"Missing {elem_key} in arrays")
                return None
            if file_metadata and zd_marker in file_metadata:
                elements.append(mx.zeros(_decode_shape(file_metadata[zd_marker])))
            else:
                elements.append(arrays[elem_key])
    else:
        # V2 polyfill: legacy ``{prefix}_keys`` / ``{prefix}_values``.
        keys_key = f"{prefix}_keys"
        values_key = f"{prefix}_values"
        if keys_key not in arrays or values_key not in arrays:
            return None
        k_zd = f"{prefix}_keys_zero_dim"
        v_zd = f"{prefix}_values_zero_dim"
        if file_metadata and k_zd in file_metadata:
            elements.append(mx.zeros(_decode_shape(file_metadata[k_zd])))
        else:
            elements.append(arrays[keys_key])
        if file_metadata and v_zd in file_metadata:
            elements.append(mx.zeros(_decode_shape(file_metadata[v_zd])))
        else:
            elements.append(arrays[values_key])
    return ("__nstate__", class_name, elements)


# --- Safetensors dtype mapping for background-thread-safe serialization ---
# These mappings enable writing safetensors files without any mx/Metal API,
# bypassing the bfloat16 limitation that blocked PR #16 v2 (numpy doesn't
# support bfloat16, but safetensors format natively does via "BF16" dtype).

_MX_TO_ST_DTYPE: dict[Any, str] = {}
_ST_TO_MX_DTYPE: dict[str, Any] = {}
_ST_DTYPE_TO_NP: dict[str, Any] = {}

if HAS_MLX:
    _MX_TO_ST_DTYPE = {
        mx.float16: "F16",
        mx.float32: "F32",
        mx.bfloat16: "BF16",
        mx.int8: "I8",
        mx.int16: "I16",
        mx.int32: "I32",
        mx.int64: "I64",
        mx.uint8: "U8",
        mx.uint16: "U16",
        mx.uint32: "U32",
        mx.uint64: "U64",
        mx.bool_: "BOOL",
    }
    _ST_TO_MX_DTYPE = {v: k for k, v in _MX_TO_ST_DTYPE.items()}

_ST_DTYPE_TO_NP = {
    "F16": np.float16,
    "F32": np.float32,
    "BF16": np.uint16,  # bfloat16 handled via uint16 view
    "I8": np.int8,
    "I16": np.int16,
    "I32": np.int32,
    "I64": np.int64,
    "U8": np.uint8,
    "U16": np.uint16,
    "U32": np.uint32,
    "U64": np.uint64,
    "BOOL": np.bool_,
}


def _extract_tensor_bytes(arr: mx.array) -> tuple[bytes, str, list[int]]:
    """Extract raw bytes from an mx.array.

    Materialize the array at this last-mile boundary before touching the
    Python buffer protocol. ``store_cache`` may create lazy block slices,
    clones, or placeholder arrays after scheduler-side pre-eval collection,
    and ``memoryview(arr)`` would otherwise trigger an implicit eval from the
    background cache-store worker thread.

    For bfloat16 arrays, uses view(uint16) since the buffer protocol does
    not support bfloat16 directly. Materialize the view as well so the raw
    buffer read never becomes an implicit MLX eval.

    Args:
        arr: MLX array to serialize.

    Returns:
        Tuple of (raw_bytes, safetensors_dtype_string, shape_list).
    """
    mx.eval(arr)
    dtype_str = _MX_TO_ST_DTYPE[arr.dtype]
    shape = list(arr.shape)
    if arr.dtype == mx.bfloat16:
        u16 = arr.view(mx.uint16)
        mx.eval(u16)
        raw = bytes(memoryview(u16))
    else:
        raw = bytes(memoryview(arr))
    return raw, dtype_str, shape


def _restore_tensor_from_bytes(
    raw: bytes, dtype_str: str, shape: list[int]
) -> mx.array:
    """Restore an mx.array from raw bytes extracted by _extract_tensor_bytes.

    No Metal API required — uses numpy as intermediary.

    Args:
        raw: Raw tensor bytes.
        dtype_str: Safetensors dtype string (e.g., "F16", "BF16").
        shape: Tensor shape as list of ints.

    Returns:
        Restored mx.array with correct dtype and shape.
    """
    np_dtype = _ST_DTYPE_TO_NP[dtype_str]
    np_arr = np.frombuffer(raw, dtype=np_dtype)
    arr = mx.array(np_arr)
    if dtype_str == "BF16":
        arr = arr.view(mx.bfloat16)
    return arr.reshape(shape)


def _write_safetensors_no_mx(
    path: str,
    tensors_raw: dict[str, tuple[bytes, str, list[int]]],
    metadata: dict[str, str] | None = None,
) -> int:
    """Write a safetensors file without any mx/Metal API calls.

    Safe to call from background threads. Produces files fully compatible
    with mx.load(path, return_metadata=True).

    The safetensors binary format:
      [8 bytes: header_size as little-endian uint64]
      [header_size bytes: JSON header]
      [remaining bytes: concatenated tensor data]

    Args:
        path: Output file path (must include .safetensors extension).
        tensors_raw: Dict of {name: (raw_bytes, dtype_str, shape)}.
        metadata: Optional string-to-string metadata dict.

    Returns:
        Total file size in bytes.
    """
    offset = 0
    header_tensors = {}
    all_data = []

    for name, (raw, dtype_str, shape) in tensors_raw.items():
        header_tensors[name] = {
            "dtype": dtype_str,
            "shape": shape,
            "data_offsets": [offset, offset + len(raw)],
        }
        all_data.append(raw)
        offset += len(raw)

    header_dict = dict(header_tensors)
    if metadata:
        header_dict["__metadata__"] = metadata

    header_json = json.dumps(header_dict, separators=(",", ":")).encode("utf-8")
    # Safetensors spec: header must be 8-byte aligned
    pad = (8 - len(header_json) % 8) % 8
    header_json += b" " * pad

    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(header_json)))
        f.write(header_json)
        for d in all_data:
            f.write(d)

    return 8 + len(header_json) + offset


def parse_size(size_str: str) -> int:
    """
    Parse a human-readable size string to bytes.

    Args:
        size_str: Size string like "100GB", "50MB", "1TB"

    Returns:
        Size in bytes.
    """
    size_str = size_str.strip().upper()

    units = {
        "B": 1,
        "KB": 1024,
        "MB": 1024**2,
        "GB": 1024**3,
        "TB": 1024**4,
    }

    for unit, multiplier in units.items():
        if size_str.endswith(unit):
            try:
                value = float(size_str[: -len(unit)])
                return int(value * multiplier)
            except ValueError:
                pass

    # Try parsing as plain number (bytes)
    try:
        return int(size_str)
    except ValueError:
        raise ValueError(f"Invalid size string: {size_str}")


@dataclass
class PagedSSDBlockMetadata:
    """
    Metadata for a block stored on SSD.

    Attributes:
        block_hash: Content hash (SHA256) for identification
        file_path: Full path to safetensors file
        file_size: Size in bytes
        token_count: Number of tokens in this block
        created_at: Timestamp when saved
        last_access: Last access time for LRU tracking
        num_layers: Number of model layers
        model_name: Model name for cache isolation between different models
        block_size: Paged cache block size that created this block
        cache_signature: Compatibility signature for the saved cache layout
        layer_cache_types: Per-layer cache type names (e.g., ["KVCache", "ArraysCache"])
        layer_meta_states: Per-layer meta_state tuples for reconstruction
    """

    block_hash: bytes
    file_path: Path
    file_size: int
    token_count: int
    created_at: float
    last_access: float
    num_layers: int
    model_name: str = ""
    block_size: int = 0
    cache_signature: str = ""
    layer_cache_types: list[str] | None = None
    layer_meta_states: list[tuple] | None = None

    def touch(self) -> None:
        """Update last access time."""
        self.last_access = time.time()

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        result = {
            "block_hash": self.block_hash.hex(),
            "file_path": str(self.file_path),
            "file_size": self.file_size,
            "token_count": self.token_count,
            "created_at": self.created_at,
            "last_access": self.last_access,
            "num_layers": self.num_layers,
            "model_name": self.model_name,
            "block_size": self.block_size,
            "cache_signature": self.cache_signature,
        }
        if self.layer_cache_types:
            result["layer_cache_types"] = self.layer_cache_types
        if self.layer_meta_states:
            # Convert tuples to lists for JSON serialization
            result["layer_meta_states"] = [list(m) for m in self.layer_meta_states]
        return result

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PagedSSDBlockMetadata:
        """Create from dictionary."""
        # Parse layer_meta_states back to tuples
        layer_meta_states = None
        if "layer_meta_states" in data and data["layer_meta_states"]:
            layer_meta_states = [tuple(m) for m in data["layer_meta_states"]]

        return cls(
            block_hash=bytes.fromhex(data["block_hash"]),
            file_path=Path(data["file_path"]),
            file_size=data["file_size"],
            token_count=data["token_count"],
            created_at=data["created_at"],
            last_access=data["last_access"],
            num_layers=data["num_layers"],
            model_name=data.get("model_name", ""),
            block_size=data.get("block_size", 0),
            cache_signature=data.get("cache_signature", ""),
            layer_cache_types=data.get("layer_cache_types"),
            layer_meta_states=layer_meta_states,
        )


class PagedSSDCacheIndex:
    """
    In-memory index of SSD cache files.

    Provides O(1) lookup by block_hash and LRU tracking for size management.
    Thread-safe for concurrent access.
    """

    def __init__(self, max_size_bytes: int):
        """
        Initialize the SSD cache index.

        Args:
            max_size_bytes: Maximum total size of SSD cache files.
        """
        self._index: dict[bytes, PagedSSDBlockMetadata] = {}
        self._lru: OrderedDict[bytes, float] = OrderedDict()
        self._total_size: int = 0
        self._max_size: int = max_size_bytes
        self._lock = threading.RLock()

    def add(self, metadata: PagedSSDBlockMetadata) -> None:
        """
        Add a block to the index.

        Args:
            metadata: Block metadata to add.
        """
        with self._lock:
            # Remove existing entry if present
            if metadata.block_hash in self._index:
                old_meta = self._index[metadata.block_hash]
                self._total_size -= old_meta.file_size
                del self._lru[metadata.block_hash]

            self._index[metadata.block_hash] = metadata
            self._lru[metadata.block_hash] = metadata.last_access
            self._total_size += metadata.file_size

    def sort_lru_by_last_access(self) -> None:
        """Restore LRU ordering from each entry's last access timestamp."""
        with self._lock:
            self._lru = OrderedDict(
                sorted(
                    (
                        (block_hash, self._index[block_hash].last_access)
                        for block_hash in self._lru
                        if block_hash in self._index
                    ),
                    key=lambda item: item[1],
                )
            )

    def get(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
        """
        Get block metadata by hash.

        Args:
            block_hash: Block content hash.

        Returns:
            PagedSSDBlockMetadata if found, None otherwise.
        """
        with self._lock:
            return self._index.get(block_hash)

    def remove(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
        """
        Remove a block from the index.

        Args:
            block_hash: Block content hash.

        Returns:
            Removed metadata if found, None otherwise.
        """
        with self._lock:
            if block_hash not in self._index:
                return None

            metadata = self._index.pop(block_hash)
            del self._lru[block_hash]
            self._total_size -= metadata.file_size
            return metadata

    def touch(self, block_hash: bytes) -> None:
        """
        Update last access time (move to end of LRU).

        Args:
            block_hash: Block content hash.
        """
        with self._lock:
            if block_hash in self._index:
                self._index[block_hash].touch()
                self._lru.move_to_end(block_hash)
                self._lru[block_hash] = self._index[block_hash].last_access

    def get_lru_entries(self, count: int) -> list[PagedSSDBlockMetadata]:
        """
        Get least recently used entries.

        Args:
            count: Maximum number of entries to return.

        Returns:
            List of LRU metadata entries.
        """
        with self._lock:
            result = []
            for block_hash in list(self._lru.keys())[:count]:
                if block_hash in self._index:
                    result.append(self._index[block_hash])
            return result

    def evict_until_size(
        self,
        target_size: int,
        max_count: int | None = None,
    ) -> list[PagedSSDBlockMetadata]:
        """
        Evict LRU entries until total size is below target.

        Args:
            target_size: Target total size in bytes.
            max_count: Optional cap on the number of entries removed in
                one call. When the cap is hit before ``total_size`` drops
                below ``target_size`` the call returns the partial slice
                and leaves the remaining LRU entries in the index; the
                caller is expected to retry on the next save. The cap is
                pushed down here (rather than the caller popping a
                surplus and reinserting it) so the index never exposes a
                transient "evicted but not yet unlinked" gap that a
                concurrent writer's ``contains()`` check could observe
                as a deleted block.

        Returns:
            List of evicted metadata (files need to be deleted by caller).

        Note:
            Loop termination depends on ``remove()`` decrementing
            ``_total_size`` for every popped entry. If a future refactor
            moves the decrement to "after the on-disk unlink succeeds",
            this loop must also gain a "skip entries already pulled this
            pass" guard or it can spin forever when unlinks fail.
        """
        with self._lock:
            evicted = []
            while self._total_size > target_size and self._lru:
                if max_count is not None and len(evicted) >= max_count:
                    break
                # Get LRU entry (first in OrderedDict)
                block_hash = next(iter(self._lru))
                metadata = self.remove(block_hash)
                if metadata:
                    evicted.append(metadata)
            return evicted

    def contains(self, block_hash: bytes) -> bool:
        """Check if block exists in index."""
        with self._lock:
            return block_hash in self._index

    @property
    def total_size(self) -> int:
        """Get total size of indexed files."""
        with self._lock:
            return self._total_size

    @property
    def max_size(self) -> int:
        """Get maximum allowed size."""
        return self._max_size

    @property
    def count(self) -> int:
        """Get number of indexed blocks."""
        with self._lock:
            return len(self._index)

    def update_file_size(self, block_hash: bytes, actual_size: int) -> None:
        """Update file size for a block after background write completes.

        Args:
            block_hash: Block content hash.
            actual_size: Actual file size in bytes.
        """
        with self._lock:
            entry = self._index.get(block_hash)
            if entry is not None:
                self._total_size += actual_size - entry.file_size
                entry.file_size = actual_size

    def get_all_hashes(self) -> list[bytes]:
        """Get all indexed block hashes."""
        with self._lock:
            return list(self._index.keys())

    def get_all_metadata(self) -> list[PagedSSDBlockMetadata]:
        """Get a snapshot of all indexed block metadata."""
        with self._lock:
            return list(self._index.values())


@dataclass
class _HotCacheBudgetEntry:
    owner: Any
    block_hash: bytes
    size_bytes: int


class SharedHotCacheBudget:
    """Process-wide byte budget for hot cache entries across cache managers."""

    def __init__(self, max_bytes: int):
        self.max_bytes = max(0, int(max_bytes))
        self._entries: OrderedDict[tuple[int, bytes], _HotCacheBudgetEntry] = (
            OrderedDict()
        )
        self._total_bytes = 0
        self._lock = threading.RLock()

    @staticmethod
    def _key(owner: Any, block_hash: bytes) -> tuple[int, bytes]:
        return (id(owner), block_hash)

    @property
    def total_bytes(self) -> int:
        with self._lock:
            return self._total_bytes

    @property
    def remaining_bytes(self) -> int:
        with self._lock:
            return max(0, self.max_bytes - self._total_bytes)

    def touch(self, owner: Any, block_hash: bytes) -> None:
        """Mark an entry as recently used in the global LRU order."""
        with self._lock:
            key = self._key(owner, block_hash)
            if key in self._entries:
                self._entries.move_to_end(key)

    def forget(self, owner: Any, block_hash: bytes) -> None:
        """Remove one entry from budget accounting if present."""
        with self._lock:
            key = self._key(owner, block_hash)
            entry = self._entries.pop(key, None)
            if entry is not None:
                self._total_bytes = max(0, self._total_bytes - entry.size_bytes)

    def forget_owner(self, owner: Any) -> None:
        """Remove all entries owned by a cache manager."""
        owner_id = id(owner)
        with self._lock:
            keys = [key for key in self._entries if key[0] == owner_id]
            for key in keys:
                entry = self._entries.pop(key)
                self._total_bytes = max(0, self._total_bytes - entry.size_bytes)

    def clear_all_owners(self) -> int:
        """Clear the hot cache of every manager the budget still references.

        The budget keeps a strong reference to each owning manager, so a
        manager orphaned by an abnormal teardown stays reachable here even
        when it is no longer attached to a loaded scheduler. Snapshot the
        owners under the lock, then clear outside it (clear_hot_cache calls
        forget_owner, which re-takes the lock).
        """
        with self._lock:
            owners = []
            seen = set()
            for entry in self._entries.values():
                if id(entry.owner) not in seen:
                    seen.add(id(entry.owner))
                    owners.append(entry.owner)
        cleared = 0
        for owner in owners:
            fn = getattr(owner, "clear_hot_cache", None)
            if callable(fn):
                try:
                    cleared += fn()
                except Exception:
                    # Keep going for the other owners, but do not hide the
                    # failure: a silently-swallowed error makes admin recovery
                    # look successful while memory is still pinned.
                    logger.warning(
                        "clear_hot_cache failed for an orphaned owner",
                        exc_info=True,
                    )
        return cleared

    def shrink_to(
        self,
        target_bytes: int,
        protected_hashes: set[bytes] | None = None,
    ) -> int:
        """Shrink the shared hot cache to ``target_bytes`` by global LRU order.

        Returns the budgeted bytes removed from hot-cache ownership. Dirty
        evictions are handed back to each owner so the existing SSD write-through
        path can preserve them.
        """
        target_bytes = max(0, int(target_bytes))
        protected_hashes = protected_hashes or set()
        victims: list[tuple[Any, bytes, int]] = []

        with self._lock:
            while self._total_bytes > target_bytes and self._entries:
                victim_key = None
                victim = None
                for key, candidate in self._entries.items():
                    if candidate.block_hash not in protected_hashes:
                        victim_key = key
                        victim = candidate
                        break
                if victim_key is None or victim is None:
                    break

                self._entries.pop(victim_key)
                self._total_bytes = max(0, self._total_bytes - victim.size_bytes)
                victims.append((victim.owner, victim.block_hash, victim.size_bytes))

        freed = 0
        for owner, block_hash, size_bytes in victims:
            evicted = owner._hot_cache_remove(block_hash, update_budget=False)
            if evicted is not None:
                freed += size_bytes
                owner._handle_hot_cache_eviction(block_hash, evicted)
        return freed

    def put(
        self, owner: Any, block_hash: bytes, size_bytes: int
    ) -> list[tuple[Any, bytes]]:
        """Account an entry and return globally-evicted owners/block hashes."""
        victims: list[tuple[Any, bytes]] = []
        size_bytes = max(0, int(size_bytes))
        with self._lock:
            key = self._key(owner, block_hash)
            old = self._entries.pop(key, None)
            if old is not None:
                self._total_bytes = max(0, self._total_bytes - old.size_bytes)

            self._entries[key] = _HotCacheBudgetEntry(
                owner=owner,
                block_hash=block_hash,
                size_bytes=size_bytes,
            )
            self._total_bytes += size_bytes

            while self._total_bytes > self.max_bytes and self._entries:
                victim_key, victim = self._entries.popitem(last=False)
                if victim_key == key and not self._entries:
                    self._entries[victim_key] = victim
                    break
                self._total_bytes = max(0, self._total_bytes - victim.size_bytes)
                victims.append((victim.owner, victim.block_hash))

        return victims


class PagedSSDCacheManager(CacheManager):
    """
    Manages SSD storage for KV cache blocks.

    Features:
    - Block-level safetensors serialization
    - Hash-based subdirectory structure (single level: /a/, /b/, etc.)
    - LRU-based SSD cache size management

    Implements the CacheManager ABC interface for consistency with other
    cache implementations in oMLX.

    Example:
        >>> manager = PagedSSDCacheManager(
        ...     cache_dir=Path("/tmp/ssd_cache"),
        ...     max_size_bytes=100 * 1024**3,  # 100GB
        ... )
        >>> manager.save_block(block_hash, cache_data, token_count=64)
        >>> loaded = manager.load_block(block_hash)
    """

    # Subdirectory prefixes (hash first char)
    SUBDIR_CHARS = "0123456789abcdef"

    def __init__(
        self,
        cache_dir: Path | None,
        max_size_bytes: int,
        hot_cache_max_bytes: int = 0,
        hot_cache_only: bool = False,
        hot_cache_budget: SharedHotCacheBudget | None = None,
        expected_model_name: str = "",
        expected_num_layers: int = 0,
        expected_block_size: int = 0,
        expected_block_size_tokens: int = _DEFAULT_BLOCK_SIZE_TOKENS,
        expected_kv_bytes_per_token: int = _DEFAULT_KV_BYTES_PER_TOKEN,
        expected_layer_cache_types: list[str] | None = None,
    ):
        """
        Initialize the SSD cache manager.

        Args:
            cache_dir: Directory for SSD cache files.
            max_size_bytes: Maximum total size of SSD cache.
            hot_cache_max_bytes: Maximum in-memory hot cache size in bytes.
                0 means disabled (default).
            hot_cache_only: When True, skip directory init and writer thread.
                All data is stored exclusively in the hot cache (RAM only).
                No SSD I/O is performed.
            hot_cache_budget: Optional process-wide hot cache budget shared
                by all loaded model cache managers.
            expected_model_name: Current model name. Blocks saved for a
                different model name are skipped at startup. Empty string
                disables this check (backwards compatible).
            expected_num_layers: Current cache-layer count. Blocks saved with
                a different num_layers are skipped at startup. 0 disables this
                check (backwards compatible). Catches stale blocks left over
                after a model upgrade changes its effective layer count (e.g.,
                #1404 attaching MTPModule changed 30 -> 40 layers).
            expected_block_size: Current paged cache block size. Blocks saved
                with another block size are skipped at startup. 0 disables this
                check for backwards compatibility.
            expected_block_size_tokens: Paged-cache block size in tokens used to
                size the pending-writes queue. Separate from ``expected_block_size``
                so the writer-queue formula keeps a real value (default 256)
                even when the cache-compat check is disabled (0). Passing a
                larger value shrinks the cap so small Macs with large blocks
                don't pin gigabytes at saturation; passing a smaller value lets
                the cap grow to give workloads with many tiny blocks enough
                burst headroom.
            expected_kv_bytes_per_token: Per-token KV byte estimate (all
                layers, K + V, dtype). Together with ``expected_block_size_tokens``
                this drives the bytes-aware queue cap. Defaults to a
                35B-class bf16 estimate; pass an explicit value for
                quantized models or unusually wide/narrow architectures.
            expected_layer_cache_types: Optional current cache layout. When
                provided, blocks with a different per-layer type list are
                skipped at startup.
        """
        self._cache_dir = cache_dir
        self._max_size = max_size_bytes
        self._index = PagedSSDCacheIndex(max_size_bytes)
        self._incompatible_index = PagedSSDCacheIndex(max_size_bytes)
        self._hot_cache_only = hot_cache_only
        self._expected_model_name = expected_model_name
        self._expected_num_layers = expected_num_layers
        self._expected_block_size = expected_block_size
        self._expected_layer_cache_types = expected_layer_cache_types
        # TurboQuant bit depth requests will quantize at; learned together
        # with the layer signature via ``set_expected_layer_signature``
        # (the depth is only known once the engine has applied the model's
        # TurboQuant settings, after this manager is constructed).
        self._expected_turboquant_kv_bits: float | None = None
        # Sub composition of mixed CacheList layers (see
        # ``cachelist_subtypes_from_cache_list``); learned together with
        # the layer signature. None disables the check (legacy managers /
        # models without mixed CacheList layers).
        self._expected_cachelist_subtypes: dict[str, list[str]] | None = None
        # Set once we have swept stale-signature blocks for the current
        # ``_expected_layer_cache_types`` / ``_expected_turboquant_kv_bits``.
        # Re-assigning the signature (e.g., via
        # ``adopt_layer_signature_if_unset``) resets this so the new
        # signature triggers its own one-shot sweep.
        self._signature_sweep_completed = False
        self._lock = threading.RLock()

        # Disk usage cache for dynamic effective max size (30s TTL)
        self._disk_usage_cache = None  # type: shutil._ntuple_diskusage | None
        self._disk_usage_cache_time: float = 0.0
        self._last_disk_pressure_warn: float = 0.0

        # Statistics
        self._stats = {
            "saves": 0,
            "saves_persisted": 0,
            "loads": 0,
            "hits": 0,
            "misses": 0,
            "evictions": 0,
            "evict_unlink_failures": 0,
            "errors": 0,
            "hot_cache_hits": 0,
            "hot_cache_evictions": 0,
            "hot_cache_promotions": 0,
            "preload_calls": 0,
            "preload_blocks_loaded": 0,
            "preload_time_ms": 0.0,
            "ssd_write_drops": 0,
            "ssd_inline_write_fallbacks": 0,
        }

        # --- Hot cache (in-memory raw-bytes tier) ---
        self._hot_cache_budget = hot_cache_budget
        self._hot_cache_max_bytes = (
            hot_cache_budget.max_bytes
            if hot_cache_budget is not None
            else hot_cache_max_bytes
        )
        self._hot_cache_enabled = self._hot_cache_max_bytes > 0
        self._hot_cache: OrderedDict[bytes, dict] = OrderedDict()
        self._hot_cache_total_bytes: int = 0
        self._hot_cache_lock = threading.Lock()

        # Initialize directory structure and scan existing files
        # Skip in hot_cache_only mode: no SSD I/O, so no directories needed.
        if self._cache_dir and not self._hot_cache_only:
            self._init_directories()
            self._scan_existing_files()

        # --- Background writer for non-blocking saves ---
        # Recompute the pending-writes cap from THIS cache's block/model
        # parameters so a non-default block size shrinks (or grows) the
        # cap appropriately. Falls back to the module-level constant
        # when no override is supplied.
        #
        # Stash the inputs the constructor was called with so callers
        # (and the plumbing-regression test) can verify what reached
        # the manager without depending on the cap math landing in a
        # particular floor/ceiling band on the test host.
        self._expected_block_size_tokens = expected_block_size_tokens
        self._expected_kv_bytes_per_token = expected_kv_bytes_per_token
        self._max_pending_writes = _compute_max_pending_writes(
            block_size_tokens=expected_block_size_tokens,
            kv_bytes_per_token=expected_kv_bytes_per_token,
        )
        self._write_queue: queue.Queue = queue.Queue(maxsize=self._max_pending_writes)
        # Track which block hashes are queued for background write
        self._pending_write_hashes: set = set()
        self._pending_write_hashes_lock = threading.Lock()
        # Lock ordering invariant: _hot_cache_lock -> _pending_write_hashes_lock.
        # Never acquire in reverse. Load path: _hot_cache_get (holds _hot_cache_lock,
        # releases), then _pending_write_buffer_get (holds _pending_write_hashes_lock).
        # Eviction path: _hot_cache_put (holds _hot_cache_lock, releases), then
        # _enqueue_ssd_write (holds _pending_write_hashes_lock).
        self._pending_write_buffers: dict[bytes, dict] = {}
        self._writer_shutdown = threading.Event()
        # Writer thread is only needed when writing to SSD.
        self._writer_thread = None
        if not self._hot_cache_only:
            self._writer_thread = threading.Thread(
                target=self._writer_loop,
                name="ssd-cache-writer",
                daemon=True,
            )
            self._writer_thread.start()

        hot_info = ""
        if self._hot_cache_enabled:
            hot_info = f", hot_cache={format_bytes(hot_cache_max_bytes)}"
        # Log initialization with disk space info
        disk_info = ""
        if self._cache_dir:
            try:
                du = shutil.disk_usage(self._cache_dir)
                disk_info = (
                    f", disk_free={format_bytes(du.free)}, "
                    f"cache_used={format_bytes(self._tracked_ssd_size())}, "
                    f"incompatible_files={self._incompatible_index.count}"
                )
            except OSError:
                pass
        logger.info(
            f"PagedSSDCacheManager initialized: dir={self._cache_dir}, "
            f"max_size={format_bytes(max_size_bytes)}{hot_info}, "
            f"existing_files={self._index.count}{disk_info}"
        )

    # --- Hot cache helpers ---

    @staticmethod
    def _hot_cache_entry_size(entry: dict) -> int:
        """Calculate memory footprint of a hot cache entry.

        Entries from save_block() use 'tensors_raw' (raw bytes).
        Entries from _promote_to_hot_cache() may use 'arrays' (mx.array objects
        loaded from SSD, not from active inference — safe to retain).
        """
        if "arrays" in entry:
            return sum(arr.nbytes for arr in entry["arrays"].values())
        if "tensors_raw" in entry:
            return sum(len(raw) for raw, _, _ in entry["tensors_raw"].values())
        return 0

    def _effective_hot_cache_max_bytes(self) -> int:
        if self._hot_cache_budget is not None:
            return self._hot_cache_budget.max_bytes
        return self._hot_cache_max_bytes

    def _hot_cache_available_bytes(self) -> int:
        if self._hot_cache_budget is not None:
            return self._hot_cache_budget.remaining_bytes
        return max(0, self._hot_cache_max_bytes - self._hot_cache_total_bytes)

    def _handle_hot_cache_eviction(self, block_hash: bytes, entry: dict) -> None:
        self._stats["hot_cache_evictions"] += 1
        if not entry.get("dirty", True):
            logger.debug(
                "Evicted clean hot cache block %s; SSD copy already exists",
                block_hash.hex()[:16],
            )
            return
        self._enqueue_ssd_write(block_hash, entry)

    def _hot_cache_put(self, block_hash: bytes, entry: dict) -> None:
        """Add entry to hot cache, evicting LRU entries if capacity exceeded.

        Evicted entries are flushed to SSD via the background writer thread.
        """
        entry_size = self._hot_cache_entry_size(entry)
        evicted_entries: list = []

        if self._hot_cache_budget is not None:
            with self._hot_cache_lock:
                old = self._hot_cache.pop(block_hash, None)
                if old is not None:
                    self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)
                self._hot_cache[block_hash] = entry
                self._hot_cache_total_bytes += entry_size

            victims = self._hot_cache_budget.put(self, block_hash, entry_size)
            for owner, victim_hash in victims:
                evicted = owner._hot_cache_remove(victim_hash, update_budget=False)
                if evicted is not None:
                    owner._handle_hot_cache_eviction(victim_hash, evicted)
            return

        with self._hot_cache_lock:
            # Remove old entry if updating
            if block_hash in self._hot_cache:
                old = self._hot_cache.pop(block_hash)
                self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)

            # Evict LRU entries until we have room
            while (
                self._hot_cache_total_bytes + entry_size > self._hot_cache_max_bytes
                and self._hot_cache
            ):
                evicted_hash, evicted = self._hot_cache.popitem(last=False)
                self._hot_cache_total_bytes -= self._hot_cache_entry_size(evicted)
                evicted_entries.append((evicted_hash, evicted))

            self._hot_cache[block_hash] = entry
            self._hot_cache_total_bytes += entry_size

        # Flush evicted entries to SSD outside the hot cache lock
        for evicted_hash, evicted in evicted_entries:
            self._handle_hot_cache_eviction(evicted_hash, evicted)

    def _enqueue_ssd_write(
        self,
        block_hash: bytes,
        entry: dict,
        *,
        blocking: bool = False,
    ) -> bool:
        """Enqueue a hot cache entry for SSD background write.

        Used when evicting from hot cache or flushing on shutdown.
        Adds block to SSD index before enqueueing write.

        All callers wait briefly for queue space. If saturation persists, the
        caller writes inline so dirty hot-cache blocks are never dropped just
        because the background writer is behind.
        """
        if self._hot_cache_only:
            return False
        if not entry.get("dirty", True):
            return True

        blk_meta = entry.get("block_metadata")
        if blk_meta is None:
            return False
        file_path = blk_meta.file_path
        tensors_raw = entry.get("tensors_raw", {})
        if not tensors_raw:
            return False
        metadata = entry["file_metadata"]

        # 1. Buffer first — instant read-back for concurrent loads (CPD K1).
        #    Must precede _index.add so load_block never sees an index hit
        #    for a block that has no file and no buffer entry yet.
        with self._pending_write_hashes_lock:
            if block_hash in self._pending_write_buffers:
                return True
            self._pending_write_buffers[block_hash] = entry
            self._pending_write_hashes.add(block_hash)

        # 2. Index second — makes the block discoverable in has_block/contains.
        if not self._index.contains(block_hash):
            self._enforce_size_limit_for_new_block(blk_meta.file_size)
            self._incompatible_index.remove(block_hash)
            self._index.add(blk_meta)

        # 3. Queue third — enqueue for background writer.
        try:
            item = (block_hash, tensors_raw, metadata, file_path)
            # Non-blocking callers (hot-cache LRU spill) also wait so a
            # transient writer backlog doesn't silently drop blocks. Blocking
            # callers (shutdown flush) use the same bounded wait.
            self._write_queue.put(item, timeout=_PENDING_WRITE_PUT_TIMEOUT_SECONDS)
            logger.debug(
                f"Evicted hot cache block to SSD write queue: "
                f"{block_hash.hex()[:16]}..."
            )
            return True
        except queue.Full:
            self._stats["ssd_inline_write_fallbacks"] += 1
            logger.warning(
                f"SSD write queue saturated (cap={self._max_pending_writes}); "
                f"writing evicted block {block_hash.hex()[:16]} inline"
            )
            ok = self._write_block_file(
                block_hash,
                tensors_raw,
                metadata,
                file_path,
                source="inline-fallback",
            )
            self._clear_pending_write(block_hash)
            return ok

    def _hot_cache_get(self, block_hash: bytes) -> dict | None:
        """Get entry from hot cache, updating LRU order. Returns None on miss."""
        with self._hot_cache_lock:
            if block_hash in self._hot_cache:
                self._hot_cache.move_to_end(block_hash)
                entry = self._hot_cache[block_hash]
            else:
                return None
        if self._hot_cache_budget is not None:
            self._hot_cache_budget.touch(self, block_hash)
        return entry

    def _pending_write_buffer_get(self, block_hash: bytes) -> dict | None:
        """Get entry from pending-write buffer. Returns None on miss."""
        with self._pending_write_hashes_lock:
            return self._pending_write_buffers.get(block_hash)

    def _hot_cache_remove(
        self, block_hash: bytes, *, update_budget: bool = True
    ) -> dict | None:
        """Remove entry from hot cache if present."""
        with self._hot_cache_lock:
            old = self._hot_cache.pop(block_hash, None)
            if old:
                self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)
        if old is not None and update_budget and self._hot_cache_budget is not None:
            self._hot_cache_budget.forget(self, block_hash)
        return old

    def _promote_to_hot_cache(
        self,
        block_hash: bytes,
        arrays: dict[str, Any],
        file_metadata: Any,
        metadata: PagedSSDBlockMetadata,
    ) -> None:
        """Promote a block loaded from SSD into the hot cache."""
        try:
            promoted_raw = {}
            for name, arr in arrays.items():
                promoted_raw[name] = _extract_tensor_bytes(arr)
            entry = {
                "tensors_raw": promoted_raw,
                "file_metadata": (
                    file_metadata if isinstance(file_metadata, dict) else {}
                ),
                "num_layers": metadata.num_layers,
                "layer_cache_types": metadata.layer_cache_types,
                "block_metadata": metadata,
                "dirty": False,
            }
            self._hot_cache_put(block_hash, entry)
            self._stats["hot_cache_promotions"] += 1
        except Exception:
            pass  # Promotion failure is non-critical

    def _init_directories(self) -> None:
        """Create cache directory structure."""
        self._cache_dir.mkdir(parents=True, exist_ok=True)

        # Create subdirectories for first hex character
        for char in self.SUBDIR_CHARS:
            subdir = self._cache_dir / char
            subdir.mkdir(exist_ok=True)

    def _get_file_path(self, block_hash: bytes) -> Path:
        """
        Get file path for a block hash.

        Uses first character of hex hash as subdirectory.

        Args:
            block_hash: Block content hash.

        Returns:
            Path to the safetensors file.
        """
        hash_hex = block_hash.hex()
        subdir = hash_hex[0]  # First character
        filename = f"{hash_hex}.safetensors"
        return self._cache_dir / subdir / filename

    def _tracked_ssd_size(self) -> int:
        """Return all SSD cache bytes tracked for this shared cache directory."""
        return self._index.total_size + self._incompatible_index.total_size

    def _tracked_ssd_count(self) -> int:
        """Return compatible plus incompatible tracked SSD cache file count."""
        return self._index.count + self._incompatible_index.count

    def _scan_existing_files(self) -> None:
        """Scan cache directory for existing files and build the compatible index.

        Only blocks compatible with the currently loaded model/layout are
        indexed. Incompatible blocks are left on disk so a shared SSD cache
        directory can safely serve multiple loaded models without one model's
        startup scan deleting another model's cache.
        """
        logger.info(f"Scanning SSD cache directory: {self._cache_dir}")

        scanned = 0
        indexed = 0
        skipped_incompatible = 0
        skipped_incompatible_bytes = 0
        errors = 0

        for subdir in self.SUBDIR_CHARS:
            subdir_path = self._cache_dir / subdir
            if not subdir_path.exists():
                continue

            for file_path in subdir_path.glob("*.safetensors"):
                scanned += 1
                try:
                    metadata = self._read_file_metadata(file_path)
                    if metadata is None:
                        continue
                    if not self._is_compatible_block(metadata):
                        skipped_incompatible += 1
                        skipped_incompatible_bytes += metadata.file_size
                        self._incompatible_index.add(metadata)
                        continue
                    self._index.add(metadata)
                    indexed += 1
                except Exception as e:
                    logger.warning(f"Failed to read {file_path}: {e}")
                    errors += 1

        self._index.sort_lru_by_last_access()
        self._incompatible_index.sort_lru_by_last_access()

        log_msg = (
            f"SSD cache scan complete: scanned={scanned}, indexed={indexed}, "
            f"errors={errors}, total_size={format_bytes(self._index.total_size)}"
        )
        if skipped_incompatible > 0:
            log_msg += (
                f", skipped_incompatible={skipped_incompatible} blocks "
                f"({format_bytes(skipped_incompatible_bytes)})"
            )
        logger.info(log_msg)

        # Startup can find a cache directory that already exceeds the shared
        # SSD budget. Converge immediately before serving requests.
        tracked_size = self._tracked_ssd_size()
        if tracked_size > 0 and tracked_size > self._get_effective_max_size():
            self._enforce_size_limit_for_new_block(0, unbounded=True)

    def _is_compatible_block(self, metadata: PagedSSDBlockMetadata) -> bool:
        """Return True when a block can be indexed for this manager."""
        if self._expected_model_name and metadata.model_name:
            if metadata.model_name != self._expected_model_name:
                return False
        if self._expected_num_layers > 0 and metadata.num_layers > 0:
            if metadata.num_layers != self._expected_num_layers:
                return False
        if self._expected_block_size > 0:
            if metadata.block_size <= 0:
                return False
            if metadata.block_size != self._expected_block_size:
                return False
        if self._expected_layer_cache_types is not None:
            if _canonicalize_layer_cache_types(
                metadata.layer_cache_types
            ) != _canonicalize_layer_cache_types(self._expected_layer_cache_types):
                return False
        if (
            self._expected_layer_cache_types is not None
            and not self._is_compatible_cache_signature(metadata)
        ):
            return False
        return True

    def _is_compatible_cache_signature(self, metadata: PagedSSDBlockMetadata) -> bool:
        """Return True when a saved cache_signature matches enabled checks."""
        if not metadata.cache_signature:
            return self._signature_bits_match("")

        try:
            payload = json.loads(metadata.cache_signature)
        except (TypeError, ValueError):
            expected_signature = (
                self._expected_cache_signature()
                if self._expected_layer_cache_types is not None
                else ""
            )
            return (
                not expected_signature or metadata.cache_signature == expected_signature
            )

        if self._expected_model_name:
            if payload.get("model_name", "") != self._expected_model_name:
                return False

        if self._expected_num_layers > 0:
            try:
                num_layers = int(payload.get("num_layers", 0) or 0)
            except (TypeError, ValueError):
                num_layers = 0
            if num_layers > 0 and num_layers != self._expected_num_layers:
                return False

        if self._expected_block_size > 0:
            try:
                block_size = int(payload.get("block_size", 0) or 0)
            except (TypeError, ValueError):
                block_size = 0
            if block_size > 0 and block_size != self._expected_block_size:
                return False

        if self._expected_layer_cache_types is not None:
            layer_cache_types = payload.get("layer_cache_types")
            if not isinstance(layer_cache_types, (list, tuple)):
                return False
            if _canonicalize_layer_cache_types(
                layer_cache_types
            ) != _canonicalize_layer_cache_types(self._expected_layer_cache_types):
                return False

        # The block must PROVE a matching CacheList sub composition — an
        # unstamped or mismatched block would rebuild a CacheList whose
        # sub arity/classes disagree with the live model.
        if (
            self._expected_cachelist_subtypes is not None
            and payload.get("cachelist_subtypes") != self._expected_cachelist_subtypes
        ):
            return False

        if not self._signature_bits_match(metadata.cache_signature):
            return False

        return True

    def _signature_bits_match(self, cache_signature: str) -> bool:
        """True when a block's recorded TurboQuant depth satisfies expectations.

        With no expected depth every block passes. With one, the block must
        PROVE a matching depth: the packed state width is
        ``ceil(head_dim * bits / 32)``, so a block written at another depth —
        or one with no recorded depth (pre-depth-stamping saves) — has an
        incompatible or unverifiable width, and restoring it poisons batch
        concatenation (#2045).
        """
        if self._expected_turboquant_kv_bits is None:
            return True
        return (
            _signature_turboquant_bits(cache_signature)
            == self._expected_turboquant_kv_bits
        )

    def is_signature_compatible(self, cache_signature: str) -> bool:
        """Per-block signature gate for restore paths that bypass the
        index scan (hot-cache / pending-write loads never pass through
        ``_is_compatible_block``). Checks the expectation-gated signature
        fields: TurboQuant depth and CacheList sub composition.
        """
        if not self._signature_bits_match(cache_signature or ""):
            return False
        if (
            self._expected_cachelist_subtypes is not None
            and _signature_cachelist_subtypes(cache_signature or "")
            != self._expected_cachelist_subtypes
        ):
            return False
        return True

    def _expected_cache_signature(self) -> str:
        if (
            not self._expected_model_name
            and self._expected_num_layers <= 0
            and self._expected_block_size <= 0
            and self._expected_layer_cache_types is None
        ):
            return ""
        return _cache_compat_signature(
            model_name=self._expected_model_name,
            num_layers=self._expected_num_layers,
            block_size=self._expected_block_size,
            layer_cache_types=self._expected_layer_cache_types,
            turboquant_kv_bits=self._expected_turboquant_kv_bits,
            cachelist_subtypes=self._expected_cachelist_subtypes,
        )

    def _read_file_metadata(self, file_path: Path) -> PagedSSDBlockMetadata | None:
        """
        Read metadata from an existing cache file.

        Args:
            file_path: Path to safetensors file.

        Returns:
            PagedSSDBlockMetadata if valid, None otherwise.
        """
        if not HAS_MLX:
            return None

        try:
            # Load just the metadata without loading tensors
            _, metadata = mx.load(str(file_path), return_metadata=True)

            block_hash_hex = metadata.get("block_hash", "")
            if not block_hash_hex:
                return None

            # Reject pre-fix blocks. RotatingKVCache layers in those files
            # may have been zero-padded to max_size, which the new merge
            # contract would treat as real attention keys. See #934 / #903
            # and the _CACHE_FORMAT_VERSION docstring for context.
            #
            # V3 polyfills V2 blocks at read time so already-stored caches
            # stay valid after the N-tuple state refactor. Versions outside
            # _READABLE_CACHE_FORMAT_VERSIONS are still rejected.
            cache_version = metadata.get("omlx_cache_format_version")
            if cache_version not in _READABLE_CACHE_FORMAT_VERSIONS:
                logger.debug(
                    "Skipping cache block with unsupported format version "
                    "%r (readable %r): %s",
                    cache_version,
                    sorted(_READABLE_CACHE_FORMAT_VERSIONS),
                    file_path,
                )
                return None

            file_stat = file_path.stat()

            # Parse cache type information if present
            layer_cache_types = None
            layer_meta_states = None

            # A present-but-unparseable field is corruption (torn write,
            # damaged file), not a legacy block: indexing the block with the
            # field silently dropped would make reconstruction guess layer
            # types or per-layer meta for its tensors. Treat the block as
            # unusable; the hash then misses and the next request re-stores
            # it. Absent fields (legacy blocks) still pass through.
            if "layer_cache_types" in metadata and metadata["layer_cache_types"]:
                try:
                    layer_cache_types = json.loads(metadata["layer_cache_types"])
                except (json.JSONDecodeError, TypeError):
                    logger.warning(
                        "Corrupt layer_cache_types metadata JSON in cache "
                        "file %s; treating the block as unusable.",
                        file_path,
                    )
                    return None

            if "layer_meta_states" in metadata and metadata["layer_meta_states"]:
                try:
                    raw_meta_states = json.loads(metadata["layer_meta_states"])
                    layer_meta_states = [tuple(m) if m else () for m in raw_meta_states]
                except (json.JSONDecodeError, TypeError):
                    logger.warning(
                        "Corrupt layer_meta_states metadata JSON in cache "
                        "file %s; treating the block as unusable.",
                        file_path,
                    )
                    return None

            return PagedSSDBlockMetadata(
                block_hash=bytes.fromhex(block_hash_hex),
                file_path=file_path,
                file_size=file_stat.st_size,
                token_count=int(metadata.get("token_count", 0)),
                created_at=file_stat.st_ctime,
                last_access=file_stat.st_mtime,
                num_layers=int(metadata.get("num_layers", 0)),
                model_name=metadata.get("model_name", ""),
                block_size=int(metadata.get("block_size", 0)),
                cache_signature=metadata.get("cache_signature", ""),
                layer_cache_types=layer_cache_types,
                layer_meta_states=layer_meta_states,
            )
        except Exception as e:
            logger.debug(f"Failed to read metadata from {file_path}: {e}")
            return None

    def _write_block_file(
        self,
        block_hash: bytes,
        tensors_raw: dict[str, Any],
        metadata: dict[str, str],
        file_path: Path,
        *,
        source: str,
    ) -> bool:
        """Write one serialized block to disk from raw tensor bytes."""
        temp_path = None
        try:
            file_path.parent.mkdir(parents=True, exist_ok=True)
            temp_path = file_path.with_name(file_path.stem + "_tmp.safetensors")
            actual_size = _write_safetensors_no_mx(
                str(temp_path), tensors_raw, metadata
            )

            os.rename(str(temp_path), str(file_path))

            # The block is now durable on disk; bump the persist counter
            # before any cleanup so ``saves_persisted`` reflects rename
            # success even if the post-rename eviction check below unlinks it.
            self._stats["saves_persisted"] += 1
            self._index.update_file_size(block_hash, actual_size)

            # Check if block was evicted while write was pending.
            if not self._index.contains(block_hash):
                logger.debug(
                    "Block %s evicted during %s write, cleaning up file",
                    block_hash.hex()[:16],
                    source,
                )
                with contextlib.suppress(Exception):
                    file_path.unlink()
            return True
        except Exception as e:
            if isinstance(e, OSError) and e.errno in (
                errno.ENOSPC,
                errno.EDQUOT,
            ):
                # Background writes may fail after save_block already returned
                # True, while inline fallbacks can still report False to the
                # caller. In both cases, surface disk pressure at ERROR level
                # and force the next save to recompute available space.
                logger.error(
                    "SSD cache disk full, cannot write block %s via %s: %s "
                    "(subsequent saves will recompute disk pressure)",
                    block_hash.hex()[:16],
                    source,
                    e,
                )
                # Invalidate the 30s disk-usage snapshot so the next
                # save sees the true (now-critical) free space and evicts
                # aggressively rather than trusting a stale inflated limit.
                # In-flight saves that already passed
                # _enforce_size_limit_for_new_block are still queued and may
                # ENOSPC again; invalidation only protects the next round of
                # save_block calls.
                with self._lock:
                    self._disk_usage_cache = None
            else:
                logger.error(
                    "SSD cache %s write failed for %s: %s",
                    source,
                    block_hash.hex()[:16],
                    e,
                )
            self._stats["errors"] += 1
            self._index.remove(block_hash)
            for p in (temp_path, file_path):
                with contextlib.suppress(Exception):
                    if p is not None and isinstance(p, Path) and p.exists():
                        p.unlink()
            return False

    def _clear_pending_write(
        self, block_hash: bytes, *, remove_hot_cache: bool = False
    ) -> None:
        """Clear pending-write bookkeeping after a queued or inline write."""
        with self._pending_write_hashes_lock:
            self._pending_write_hashes.discard(block_hash)
            self._pending_write_buffers.pop(block_hash, None)
        if remove_hot_cache:
            self._hot_cache_remove(block_hash)

    def _writer_loop(self) -> None:
        """Background writer that drains the write queue.

        Runs in a dedicated daemon thread. Writes full safetensors files
        using pure Python I/O (no mx/Metal API calls), then atomically
        renames temp files to their final paths.

        This is safe because save_block() extracts tensor data as raw bytes
        on the inference thread (Metal-safe), and this thread only performs
        standard file I/O operations.
        """
        while True:
            item = None
            try:
                item = self._write_queue.get(timeout=1.0)
            except queue.Empty:
                # Exit if shutdown was requested and queue is empty
                if self._writer_shutdown.is_set():
                    break
                continue

            if item is None:  # Sentinel for shutdown
                break

            block_hash, tensors_raw, metadata, file_path = item
            try:
                self._write_block_file(
                    block_hash, tensors_raw, metadata, file_path, source="background"
                )
                self._clear_pending_write(
                    block_hash, remove_hot_cache=not self._hot_cache_enabled
                )
            finally:
                # Avoid pinning the last raw tensor-byte batch while the
                # writer thread blocks waiting for more work.
                item = None
                block_hash = tensors_raw = metadata = file_path = None

    def save_block(
        self,
        block_hash: bytes,
        cache_data: list[Any],
        token_count: int,
        model_name: str = "",
        layer_cache_types: list[str] | None = None,
        layer_meta_states: list[tuple] | None = None,
        hot_cache_write_back: bool = True,
    ) -> bool:
        """
        Save a KV cache block to SSD storage (non-blocking).

        Data is enqueued for background writing. The block is immediately
        available for reads via the in-memory pending-writes buffer.

        Args:
            block_hash: Content hash for the block.
            cache_data: List of per-layer data. Each element is either:
                - (keys, values) tuple for standard caches (KVCache, etc.)
                - ('__cache_list__', sub_tensors) marker tuple for CacheList layers,
                  where sub_tensors is List[Tuple[keys, values]] per sub-cache.
            token_count: Number of tokens in the block.
            model_name: Model name for cache isolation between different models.
            layer_cache_types: Optional list of cache type names per layer
                (e.g., ["KVCache", "ArraysCache", "KVCache", "CacheList"]).
            layer_meta_states: Optional list of meta_state tuples per layer
                for reconstruction (e.g., [(offset,), (keep, max_size, offset, _idx)]).
            hot_cache_write_back: When False in SSD-backed hot-cache mode, enqueue
                through the SSD writer path instead of retaining a hot-cache copy.

        Returns:
            True if enqueued successfully, False otherwise.
        """
        if not HAS_MLX:
            logger.error("MLX not available, cannot save block")
            return False

        layer_cache_types = _storage_layer_cache_types(layer_cache_types)

        # First save call after a model load is the canonical source for
        # the live layer-cache signature (post-TurboQuant / post-MTP). If
        # the manager wasn't told the signature at construction, adopt it
        # now and sweep any index entries left over from a prior config.
        if self.adopt_layer_signature_if_unset(layer_cache_types):
            try:
                self.invalidate_stale_layer_signature()
            except Exception as e:
                logger.warning("Stale-signature sweep failed: %s", e)

        # Check if already exists in index (thread-safe)
        if self._index.contains(block_hash):
            self._index.touch(block_hash)
            self._stats["hits"] += 1
            return True

        # Also check hot cache / pending writes buffer
        hot_entry = None
        with self._hot_cache_lock:
            if block_hash in self._hot_cache:
                hot_entry = self._hot_cache[block_hash]

        if hot_entry is not None:
            if hot_cache_write_back or self._hot_cache_only:
                self._stats["hits"] += 1
                return True
            if not hot_entry.get("dirty", True):
                self._stats["hits"] += 1
                return True
            # Pressure write-through for an already-dirty hot-cache entry:
            # use the same pending-buffer / SSD-writer path as hot-cache
            # eviction, then drop the long-lived hot-cache reference.
            if self._enqueue_ssd_write(block_hash, hot_entry):
                self._hot_cache_remove(block_hash)
                self._stats["hits"] += 1
                return True
            return False

        file_path = self._get_file_path(block_hash)

        try:

            # Prepare arrays for safetensors. Three layer_data shapes are
            # accepted:
            # - ``('__nstate__', class_name, [elem0, elem1, ...])`` — V3
            #   N-tuple state from a handler-driven serialize_state path.
            # - ``('__cache_list__', sub_tensors)`` — composite layer; each
            #   sub_tensor may itself be a 2-tuple ``(keys, values)`` (V2
            #   legacy from prefix_cache) or an ``__nstate__`` marker.
            # - ``('__turboquant__'/'__turboquant_v2__', ...)`` — bespoke
            #   TurboQuant payload, unchanged.
            # - ``(keys, values)`` 2-tuple — V2 legacy. Promoted to V3 by
            #   storing as a length-2 ``__nstate__`` so the on-disk shape
            #   is uniform regardless of whether the producer (prefix_cache,
            #   etc.) has been migrated to emit ``__nstate__`` markers yet.
            arrays = {}
            has_pooling_cache_delta = False
            cache_list_meta = (
                {}
            )  # Per-layer sidecar metadata (sub_count, state_count, etc.)

            # Shim; module-level to avoid a recursive-closure refcount
            # cycle pinning `arrays` — see _store_nstate_elements_flat.
            def _store_nstate_elements(prefix: str, elements):
                _store_nstate_elements_flat(arrays, cache_list_meta, prefix, elements)

            for i, layer_data in enumerate(cache_data):
                if (
                    isinstance(layer_data, tuple)
                    and len(layer_data) >= 2
                    and isinstance(layer_data[0], str)
                    and layer_data[0] == "__nstate__"
                ):
                    # ('__nstate__', class_name, [elements]) — V3 native
                    class_name = layer_data[1] if len(layer_data) >= 2 else None
                    elements = layer_data[2] if len(layer_data) >= 3 else []
                    if class_name:
                        cache_list_meta[f"layer_{i}_state_class_name"] = class_name
                    if class_name == POOLING_CACHE_DELTA_CLASS:
                        has_pooling_cache_delta = True
                    _store_nstate_elements(f"layer_{i}", elements)
                elif (
                    isinstance(layer_data, tuple)
                    and len(layer_data) == 2
                    and isinstance(layer_data[0], str)
                    and layer_data[0] == "__cache_list__"
                ):
                    # CacheList: sub-indexed tensors. Each sub_tensor may be
                    # a 2-tuple (legacy) or an ``__nstate__`` marker.
                    sub_tensors = layer_data[1]
                    cache_list_meta[f"layer_{i}_sub_count"] = str(len(sub_tensors))
                    for j, sub_tensor in enumerate(sub_tensors):
                        sub_prefix = f"layer_{i}_sub_{j}"
                        if (
                            isinstance(sub_tensor, tuple)
                            and len(sub_tensor) >= 2
                            and isinstance(sub_tensor[0], str)
                            and sub_tensor[0] == "__nstate__"
                        ):
                            sub_class_name = (
                                sub_tensor[1] if len(sub_tensor) >= 2 else None
                            )
                            sub_elements = sub_tensor[2] if len(sub_tensor) >= 3 else []
                            if sub_class_name:
                                cache_list_meta[f"{sub_prefix}_state_class_name"] = (
                                    sub_class_name
                                )
                            if sub_class_name == POOLING_CACHE_DELTA_CLASS:
                                has_pooling_cache_delta = True
                            _store_nstate_elements(sub_prefix, sub_elements)
                        elif (
                            isinstance(sub_tensor, (list, tuple))
                            and len(sub_tensor) >= 2
                        ):
                            # V2 legacy: treat as N-tuple with no class name.
                            _store_nstate_elements(sub_prefix, list(sub_tensor))
                        else:
                            logger.error(
                                f"Unsupported sub_tensor format at layer {i} "
                                f"sub {j}: {type(sub_tensor).__name__}"
                            )
                            return False
                elif (
                    isinstance(layer_data, tuple)
                    and len(layer_data) == 2
                    and isinstance(layer_data[0], str)
                    and layer_data[0] in ("__turboquant__", "__turboquant_v2__")
                ):
                    # TurboQuant v2: NamedTuple states (ks, vs)
                    ks, vs = layer_data[1]
                    # Flatten NamedTuple fields into individual tensors
                    tq_tensor_idx = 0
                    for prefix, state in [("k", ks), ("v", vs)]:
                        for field_name in state._fields:
                            val = getattr(state, field_name)
                            if isinstance(val, mx.array):
                                arrays[f"layer_{i}_tq_{prefix}_{field_name}"] = val
                                tq_tensor_idx += 1
                    cache_list_meta[f"layer_{i}_turboquant_v2"] = "1"
                    cache_list_meta[f"layer_{i}_tq_key_type"] = type(ks).__name__
                    cache_list_meta[f"layer_{i}_tq_value_type"] = type(vs).__name__
                    cache_list_meta[f"layer_{i}_tq_key_fields"] = ",".join(ks._fields)
                    cache_list_meta[f"layer_{i}_tq_value_fields"] = ",".join(vs._fields)
                else:
                    # V2 legacy: 2-tuple (keys, values). Upgrade to V3
                    # __nstate__ on disk so all readers see a uniform shape.
                    if not (
                        isinstance(layer_data, (list, tuple)) and len(layer_data) >= 2
                    ):
                        logger.error(
                            f"Unsupported layer_data format at layer {i}: "
                            f"{type(layer_data).__name__}"
                        )
                        return False
                    _store_nstate_elements(f"layer_{i}", list(layer_data))

            block_size = self._expected_block_size or token_count
            # Stamp the depth the block's own TurboQuant layers were packed
            # at (observation), falling back to the manager's expectation
            # only when the block carries no meta_state. A signature must
            # never vouch for a width the payload does not have.
            block_bits = _block_turboquant_bits(layer_cache_types, layer_meta_states)
            cache_signature = _cache_compat_signature(
                model_name=model_name,
                num_layers=len(cache_data),
                block_size=block_size,
                layer_cache_types=layer_cache_types,
                turboquant_kv_bits=(
                    block_bits
                    if block_bits is not None
                    else self._expected_turboquant_kv_bits
                ),
                # Stamped from the block's own payload (observation), like
                # the TurboQuant depth above.
                cachelist_subtypes=_block_cachelist_subtypes(
                    cache_data, layer_cache_types, layer_meta_states
                ),
            )

            # Prepare metadata
            metadata = {
                "omlx_cache_format_version": (
                    POOLING_CACHE_DELTA_FORMAT_VERSION
                    if has_pooling_cache_delta
                    else _CACHE_FORMAT_VERSION
                ),
                "block_hash": block_hash.hex(),
                "token_count": str(token_count),
                "num_layers": str(len(cache_data)),
                "model_name": model_name,
                "block_size": str(block_size),
                "cache_signature": cache_signature,
                "created_at": str(time.time()),
            }

            # Add cache type information if provided
            if layer_cache_types:
                metadata["layer_cache_types"] = json.dumps(layer_cache_types)
            if layer_meta_states:
                clamped_meta_states = _clamp_rotating_meta_states(
                    cache_data, layer_cache_types, layer_meta_states
                )
                metadata["layer_meta_states"] = json.dumps(
                    [list(m) if m else [] for m in clamped_meta_states]
                )

            # Merge CacheList sub_count metadata
            metadata.update(cache_list_meta)

            # Last-mile materialization happens in _extract_tensor_bytes.
            # scheduler._cleanup_finished still pre-dispatches real KV arrays,
            # but store_cache creates additional lazy slices, clones, and
            # placeholders here after that collection step. Evaluate those
            # derived arrays before memoryview() so the buffer protocol never
            # becomes the first MLX eval site on the store-cache worker thread.
            # Race history: #978/#1040/#1106/#1437/#1558.
            tensors_raw = {}
            for name, arr in arrays.items():
                tensors_raw[name] = _extract_tensor_bytes(arr)

            # Estimate file size: raw tensor bytes + safetensors header.
            # The header is JSON-encoded per tensor (name + dtype + shape +
            # data_offsets, typically ~85 bytes) plus an 8-byte length prefix
            # and the user metadata block. Compute the metadata-JSON length
            # exactly (large `layer_meta_states` JSON on deep-layer models
            # can exceed a fixed 1 KiB constant) and keep 128 B/tensor as an
            # upper bound on the per-tensor header. The 256 B margin covers
            # the JSON separators / `__metadata__` key envelope safetensors
            # adds at write time.
            try:
                metadata_json_len = len(json.dumps(metadata).encode("utf-8"))
            except (TypeError, ValueError):
                metadata_json_len = 1024
            header_overhead = metadata_json_len + 256 + 128 * len(tensors_raw)
            estimated_size = (
                sum(len(raw) for raw, _, _ in tensors_raw.values()) + header_overhead
            )

            now = time.time()
            block_metadata = PagedSSDBlockMetadata(
                block_hash=block_hash,
                file_path=file_path,
                file_size=estimated_size,
                token_count=token_count,
                created_at=now,
                last_access=now,
                num_layers=len(cache_data),
                model_name=model_name,
                block_size=block_size,
                cache_signature=cache_signature,
                layer_cache_types=layer_cache_types,
                layer_meta_states=layer_meta_states,
            )

            # Store in hot cache (or temporary buffer) for immediate read-back.
            # Uses raw bytes (not mx.array objects) so Metal GPU memory can be
            # released as soon as the inference thread is done with the arrays.
            # NOTE: _promote_to_hot_cache() stores mx.array objects directly
            # because those are freshly loaded from SSD (not active inference),
            # so they don't tie up Metal allocations from the inference pipeline.
            # Storing live inference arrays here would accumulate GPU memory
            # under a large hot cache and cause kernel panics (IOGPUMemory underflow).
            cache_entry = {
                "tensors_raw": tensors_raw,
                "file_metadata": metadata,
                "num_layers": len(cache_data),
                "layer_cache_types": layer_cache_types,
                "block_metadata": block_metadata,
                "dirty": True,
            }

            if self._hot_cache_enabled and (
                hot_cache_write_back or self._hot_cache_only
            ):
                # Write-back mode: store only in hot cache, no SSD index entry.
                # SSD index entry is created later when block is evicted or
                # flushed to SSD (in _enqueue_ssd_write).
                self._hot_cache_put(block_hash, cache_entry)
                self._stats["saves"] += 1
                return True

            if self._hot_cache_only:
                # Hot cache disabled but hot_cache_only set: block is not retained.
                return False

            if self._hot_cache_enabled and not hot_cache_write_back:
                # Pressure write-through: keep the dirty-block durability path
                # but avoid retaining this block as a hot-cache entry.
                ok = self._enqueue_ssd_write(block_hash, cache_entry)
                if ok:
                    self._stats["saves"] += 1
                    logger.debug(
                        f"Enqueued block for SSD write-through: "
                        f"{block_hash.hex()[:16]}..., "
                        f"size={format_bytes(estimated_size)}"
                    )
                return ok

            # Evict LRU blocks to make room for the new block. Done here
            # (post-tensor-build) so the actual block size is known and the
            # cache doesn't oscillate around the configured limit.
            self._enforce_size_limit_for_new_block(estimated_size)

            # SSD path: add to index for SSD file tracking
            self._incompatible_index.remove(block_hash)
            self._index.add(block_metadata)

            # Hot cache disabled: use temporary buffer + immediate SSD write
            with self._hot_cache_lock:
                self._hot_cache[block_hash] = cache_entry

            # Track pending write
            with self._pending_write_hashes_lock:
                self._pending_write_hashes.add(block_hash)

            # Enqueue full file write for background thread. Wait on Full so
            # transient bursts (faster than the writer can drain) don't
            # immediately punch holes in the cache chain.
            try:
                self._write_queue.put(
                    (block_hash, tensors_raw, metadata, file_path),
                    timeout=_PENDING_WRITE_PUT_TIMEOUT_SECONDS,
                )
            except queue.Full:
                self._stats["ssd_inline_write_fallbacks"] += 1
                logger.warning(
                    f"SSD cache write queue saturated (cap={self._max_pending_writes}); "
                    f"writing {block_hash.hex()[:16]} inline"
                )
                ok = self._write_block_file(
                    block_hash,
                    tensors_raw,
                    metadata,
                    file_path,
                    source="inline-fallback",
                )
                self._clear_pending_write(block_hash, remove_hot_cache=True)
                if not ok:
                    return False
                self._stats["saves"] += 1
                return True

            self._stats["saves"] += 1
            logger.debug(
                f"Enqueued block for SSD cache write: {block_hash.hex()[:16]}..., "
                f"size={format_bytes(estimated_size)}"
            )
            return True

        except Exception as e:
            logger.error(f"Failed to prepare block for SSD cache: {e}")
            self._stats["errors"] += 1
            return False

    def _reconstruct_cache_data(
        self,
        arrays: dict[str, Any],
        file_metadata: dict[str, str],
        num_layers: int,
        layer_cache_types: list[str] | None = None,
    ) -> list[Any] | None:
        """Reconstruct cache_data list from flattened arrays and metadata.

        Shared helper for load_block(), load_block_with_metadata(), and
        pending-writes read path to avoid code duplication.

        Returns layer_data as one of:
        - ``('__nstate__', class_name, [elem0, elem1, ...])`` — V3 N-tuple.
        - ``('__cache_list__', sub_tensors)`` where each sub_tensor is an
          ``__nstate__`` marker — composite layer.
        - ``('__turboquant_v2__', (ks, vs))`` — TurboQuant payload (unchanged).

        V2 blocks (`layer_{i}_keys` / `layer_{i}_values` keys, no
        ``state_count`` metadata) are read via a polyfill that converts
        them to ``__nstate__`` markers with two elements, so downstream
        code paths see a uniform shape.

        Args:
            arrays: Flattened tensor dict.
            file_metadata: Safetensors metadata dict (string values).
            num_layers: Number of model layers.
            layer_cache_types: Per-layer cache type names.

        Returns:
            Reconstructed cache_data list, or None on error.
        """
        cache_data: list[Any] = []

        # When the on-disk state has exactly two elements (which covers all
        # legacy 2-tuple caches: KVCache, RotatingKVCache, ConcatenateKVCache,
        # ChunkedKVCache, QuantizedKVCache when stored as keys/values), the
        # reconstructed layer is unwrapped to a plain ``(keys, values)``
        # 2-tuple so existing callers (prefix_cache, scheduler, tests) see
        # no change. Real N-tuple caches (PoolingCache, BatchKVCache, ...)
        # surface as ``('__nstate__', class_name, elements)`` markers that
        # downstream code must dispatch on.
        def _maybe_unwrap_legacy(marker: tuple) -> Any:
            _, _, elements = marker
            # Only the legacy plain ``(keys, values)`` shape unwraps. If
            # either element is itself composite (a recursed sub-state), the
            # ``__nstate__`` wrapper must survive so the shape matches save.
            if len(elements) == 2 and not any(
                isinstance(e, (tuple, list)) for e in elements
            ):
                return (elements[0], elements[1])
            return marker

        # Shim; module-level to avoid a recursive-closure refcount
        # cycle pinning `arrays` — see _load_nstate_flat.
        def _load_nstate(prefix: str, fallback_class: str | None) -> tuple | None:
            return _load_nstate_flat(arrays, file_metadata, prefix, fallback_class)

        for i in range(num_layers):
            cache_type = (
                layer_cache_types[i]
                if layer_cache_types and i < len(layer_cache_types)
                else None
            )

            if cache_type == "CacheList":
                sub_count_key = f"layer_{i}_sub_count"
                sub_count = 0
                if file_metadata and sub_count_key in file_metadata:
                    try:
                        sub_count = int(file_metadata[sub_count_key])
                    except (ValueError, TypeError):
                        pass

                if sub_count > 0:
                    sub_tensors: list[Any] = []
                    for j in range(sub_count):
                        sub_marker = _load_nstate(
                            f"layer_{i}_sub_{j}", fallback_class=None
                        )
                        if sub_marker is None:
                            logger.error(
                                f"Missing sub-cache {j} for CacheList layer {i}"
                            )
                            return None
                        # Length-2 sub-states unwrap to (keys, values); longer
                        # N-tuples surface as ``__nstate__`` markers downstream.
                        sub_tensors.append(_maybe_unwrap_legacy(sub_marker))
                    # Preserve the legacy list shape — callers (prefix_cache,
                    # tests) expect ``cache_data[i]`` to be a list of
                    # sub-cache states for CacheList layers, not a wrapper
                    # marker.
                    cache_data.append(sub_tensors)
                else:
                    layer_marker = _load_nstate(f"layer_{i}", fallback_class=cache_type)
                    if layer_marker is None:
                        logger.error(f"Missing N-tuple state for layer {i}")
                        return None
                    cache_data.append(_maybe_unwrap_legacy(layer_marker))
            elif file_metadata and f"layer_{i}_turboquant_v2" in file_metadata:
                # TurboQuant v2: reconstruct NamedTuple states from flattened tensors
                from ..turboquant_kv import (
                    TurboQuantMSEState,
                    TurboQuantPolarProdState,
                    TurboQuantPolarState,
                    TurboQuantProdState,
                    TurboQuantSplitState,
                )

                key_type = file_metadata.get(f"layer_{i}_tq_key_type", "")
                value_type = file_metadata.get(f"layer_{i}_tq_value_type", "")
                key_fields = file_metadata.get(f"layer_{i}_tq_key_fields", "").split(
                    ","
                )
                value_fields = file_metadata.get(
                    f"layer_{i}_tq_value_fields", ""
                ).split(",")
                _type_map = {
                    "TurboQuantMSEState": TurboQuantMSEState,
                    "TurboQuantProdState": TurboQuantProdState,
                    "TurboQuantPolarState": TurboQuantPolarState,
                    "TurboQuantPolarProdState": TurboQuantPolarProdState,
                    "TurboQuantSplitState": TurboQuantSplitState,
                }
                try:
                    k_cls = _type_map[key_type]
                    v_cls = _type_map[value_type]
                    k_tensors = [arrays[f"layer_{i}_tq_k_{f}"] for f in key_fields]
                    v_tensors = [arrays[f"layer_{i}_tq_v_{f}"] for f in value_fields]
                    ks = k_cls(*k_tensors)
                    vs = v_cls(*v_tensors)
                    cache_data.append(("__turboquant_v2__", (ks, vs)))
                except (KeyError, TypeError) as e:
                    logger.error(f"TurboQuant v2 layer {i}: reconstruction failed: {e}")
                    return None
            else:
                # Standard cache layer (KVCache, RotatingKVCache,
                # PoolingCache, ...). V3 stores all state elements as
                # ``layer_{i}_state_{k}``; V2 polyfill reads the legacy
                # ``layer_{i}_keys`` / ``layer_{i}_values`` 2-tuple shape.
                # Length-2 markers unwrap to ``(keys, values)`` for legacy
                # caller compatibility; longer N-tuples (PoolingCache etc.)
                # propagate as ``__nstate__`` markers.
                layer_marker = _load_nstate(f"layer_{i}", fallback_class=cache_type)
                if layer_marker is None:
                    logger.error(f"Missing N-tuple state for layer {i}")
                    return None
                cache_data.append(_maybe_unwrap_legacy(layer_marker))

        return cache_data

    @staticmethod
    def _arrays_from_tensors_raw(
        tensors_raw: dict[str, tuple[bytes, str, list[int]]],
    ) -> dict[str, mx.array]:
        """Convert raw bytes dict back to mx.array dict for _reconstruct_cache_data.

        Args:
            tensors_raw: Dict of {name: (raw_bytes, dtype_str, shape)}.

        Returns:
            Dict of {name: mx.array} with correct dtypes and shapes.
        """
        arrays = {}
        for name, (raw, dtype_str, shape) in tensors_raw.items():
            arrays[name] = _restore_tensor_from_bytes(raw, dtype_str, shape)
        return arrays

    def load_block(self, block_hash: bytes) -> list[Any] | None:
        """
        Load a KV cache block from SSD storage.

        Checks pending writes first (in-memory, no I/O), then falls back to disk
        read with a timeout to prevent inference deadlocks.

        Args:
            block_hash: Content hash for the block.

        Returns:
            List of per-layer data, or None if not found/timed out.
            Each element is either:
            - (keys, values) tuple for standard caches
            - List[Tuple[keys, values]] for CacheList layers
        """
        if not HAS_MLX:
            logger.error("MLX not available, cannot load block")
            return None

        # Check hot cache first (in-memory, no I/O)
        entry = self._hot_cache_get(block_hash)
        if entry is not None:
            # Entries from _promote_to_hot_cache() store mx.array objects directly
            # (safe — they come from SSD loads, not active inference).
            # Entries from save_block() use tensors_raw (raw bytes).
            arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
                entry["tensors_raw"]
            )
            cache_data = self._reconstruct_cache_data(
                arrays,
                entry["file_metadata"],
                entry["num_layers"],
                entry["layer_cache_types"],
            )
            if cache_data is not None:
                self._index.touch(block_hash)
                self._stats["loads"] += 1
                self._stats["hits"] += 1
                self._stats["hot_cache_hits"] += 1
                logger.debug(f"Loaded block from hot cache: {block_hash.hex()[:16]}...")
            return cache_data

        # Check pending-write buffer (evicted from hot cache, SSD write in progress)
        entry = self._pending_write_buffer_get(block_hash)
        if entry is not None:
            arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
                entry["tensors_raw"]
            )
            cache_data = self._reconstruct_cache_data(
                arrays,
                entry["file_metadata"],
                entry["num_layers"],
                entry["layer_cache_types"],
            )
            if cache_data is not None:
                self._index.touch(block_hash)
                self._stats["loads"] += 1
                self._stats["hits"] += 1
                self._stats["hot_cache_hits"] += 1
                logger.debug(
                    f"Loaded block from pending write buffer: "
                    f"{block_hash.hex()[:16]}..."
                )
            return cache_data

        # Check index
        metadata = self._index.get(block_hash)
        if metadata is None:
            self._stats["misses"] += 1
            return None

        file_path = metadata.file_path

        if not file_path.exists():
            logger.warning(f"SSD cache file missing: {file_path}")
            self._index.remove(block_hash)
            self._stats["misses"] += 1
            return None

        try:
            # Load directly on the inference thread (Metal-safe).
            # SSD read for a ~10MB block takes ~2ms @ 5GB/s — negligible.
            # Previous executor-based approach caused deadlocks when
            # mx.load() in a worker thread contested Metal GPU resources
            # with the main inference thread.
            try:
                arrays, file_metadata = mx.load(str(file_path), return_metadata=True)
            except FileNotFoundError:
                # Concurrent evictor unlinked the file between the
                # exists() check above and this load. Treat as a miss
                # and prune the stale index entry.
                self._index.remove(block_hash)
                self._stats["misses"] += 1
                return None

            # Defensive: even if the index is stale (e.g. from a previous
            # run that pre-dates the format version field), reject blocks
            # without a readable version marker before they can poison
            # the hot cache or downstream merge logic.
            if (
                file_metadata
                and file_metadata.get("omlx_cache_format_version")
                not in _READABLE_CACHE_FORMAT_VERSIONS
            ):
                self._index.remove(block_hash)
                self._stats["misses"] += 1
                return None

            # Get layer_cache_types for CacheList detection
            layer_cache_types = metadata.layer_cache_types
            if (
                not layer_cache_types
                and file_metadata
                and "layer_cache_types" in file_metadata
            ):
                try:
                    layer_cache_types = json.loads(file_metadata["layer_cache_types"])
                except (json.JSONDecodeError, TypeError):
                    layer_cache_types = None

            cache_data = self._reconstruct_cache_data(
                arrays,
                file_metadata,
                metadata.num_layers,
                layer_cache_types,
            )
            if cache_data is None:
                return None

            # Update access time
            self._index.touch(block_hash)
            self._stats["loads"] += 1
            self._stats["hits"] += 1

            # Promote to hot cache for faster access next time
            if self._hot_cache_enabled:
                self._promote_to_hot_cache(block_hash, arrays, file_metadata, metadata)

            logger.debug(f"Loaded block from SSD cache: {block_hash.hex()[:16]}...")
            return cache_data

        except Exception as e:
            logger.error(f"Failed to load block from SSD cache: {e}")
            self._stats["errors"] += 1
            # Remove corrupted entry
            self._index.remove(block_hash)
            try:
                file_path.unlink()
            except Exception:
                pass
            return None

    def load_block_with_metadata(
        self,
        block_hash: bytes,
        promote_to_hot_cache: bool = True,
    ) -> tuple[list[Any] | None, dict[str, Any] | None]:
        """
        Load a KV cache block with its metadata from SSD storage.

        Checks pending writes first (zero I/O), then falls back to disk
        read with a timeout to prevent inference deadlocks.

        Args:
            block_hash: Content hash for the block.
            promote_to_hot_cache: When False, do not retain SSD-loaded data in
                the hot cache after reconstructing it for the active request.

        Returns:
            Tuple of (cache_data, metadata_dict) where:
            - cache_data: List of per-layer data, or None.
              Each element is either (keys, values) or List[Tuple[keys, values]]
              for CacheList layers.
            - metadata_dict: Dictionary with cache type info, or None
              {
                  "layer_cache_types": List[str],  # per-layer type names
                  "layer_meta_states": List[Tuple],  # per-layer meta states
                  "num_layers": int,
                  "token_count": int,
              }
        """
        if not HAS_MLX:
            logger.error("MLX not available, cannot load block")
            return None, None

        # Check hot cache first (in-memory, no I/O)
        entry = self._hot_cache_get(block_hash)
        if entry is not None:
            blk_meta = entry["block_metadata"]
            arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
                entry["tensors_raw"]
            )
            cache_data = self._reconstruct_cache_data(
                arrays,
                entry["file_metadata"],
                entry["num_layers"],
                entry["layer_cache_types"],
            )
            if cache_data is None:
                return None, None

            metadata_dict = {
                "num_layers": entry["num_layers"],
                "token_count": blk_meta.token_count,
                "model_name": blk_meta.model_name,
                "block_size": blk_meta.block_size,
                "cache_signature": blk_meta.cache_signature,
                "layer_cache_types": entry["layer_cache_types"],
                "layer_meta_states": blk_meta.layer_meta_states,
            }

            self._index.touch(block_hash)
            self._stats["loads"] += 1
            self._stats["hits"] += 1
            self._stats["hot_cache_hits"] += 1
            logger.debug(
                f"Loaded block with metadata from hot cache: "
                f"{block_hash.hex()[:16]}..."
            )
            return cache_data, metadata_dict

        # Check pending-write buffer (evicted from hot cache, SSD write in progress)
        entry = self._pending_write_buffer_get(block_hash)
        if entry is not None:
            blk_meta = entry["block_metadata"]
            arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
                entry["tensors_raw"]
            )
            cache_data = self._reconstruct_cache_data(
                arrays,
                entry["file_metadata"],
                entry["num_layers"],
                entry["layer_cache_types"],
            )
            if cache_data is None:
                return None, None

            metadata_dict = {
                "num_layers": entry["num_layers"],
                "token_count": blk_meta.token_count,
                "model_name": blk_meta.model_name,
                "block_size": blk_meta.block_size,
                "cache_signature": blk_meta.cache_signature,
                "layer_cache_types": entry["layer_cache_types"],
                "layer_meta_states": blk_meta.layer_meta_states,
            }

            self._index.touch(block_hash)
            self._stats["loads"] += 1
            self._stats["hits"] += 1
            self._stats["hot_cache_hits"] += 1
            logger.debug(
                f"Loaded block with metadata from pending write buffer: "
                f"{block_hash.hex()[:16]}..."
            )
            return cache_data, metadata_dict

        # Check index
        block_metadata = self._index.get(block_hash)
        if block_metadata is None:
            self._stats["misses"] += 1
            return None, None

        file_path = block_metadata.file_path

        if not file_path.exists():
            logger.warning(f"SSD cache file missing: {file_path}")
            self._index.remove(block_hash)
            self._stats["misses"] += 1
            return None, None

        try:
            # Load directly on the inference thread (Metal-safe).
            # See load_block() for rationale on removing the executor.
            arrays, file_metadata = mx.load(str(file_path), return_metadata=True)

            # Defensive version check, mirrors load_block().
            if (
                file_metadata
                and file_metadata.get("omlx_cache_format_version")
                not in _READABLE_CACHE_FORMAT_VERSIONS
            ):
                self._index.remove(block_hash)
                self._stats["misses"] += 1
                return None, None

            # Parse layer_cache_types early for CacheList detection
            layer_cache_types = block_metadata.layer_cache_types
            if (
                not layer_cache_types
                and file_metadata
                and "layer_cache_types" in file_metadata
            ):
                try:
                    layer_cache_types = json.loads(file_metadata["layer_cache_types"])
                except (json.JSONDecodeError, TypeError):
                    layer_cache_types = None

            cache_data = self._reconstruct_cache_data(
                arrays,
                file_metadata,
                block_metadata.num_layers,
                layer_cache_types,
            )
            if cache_data is None:
                return None, None

            # Build metadata dict for reconstruction
            metadata_dict = {
                "num_layers": block_metadata.num_layers,
                "token_count": block_metadata.token_count,
                "model_name": block_metadata.model_name,
                "block_size": block_metadata.block_size,
                "cache_signature": block_metadata.cache_signature,
                "layer_cache_types": layer_cache_types,
                "layer_meta_states": block_metadata.layer_meta_states,
            }

            if not metadata_dict["layer_meta_states"] and file_metadata:
                if "layer_meta_states" in file_metadata:
                    try:
                        raw = json.loads(file_metadata["layer_meta_states"])
                        metadata_dict["layer_meta_states"] = [
                            tuple(m) if m else () for m in raw
                        ]
                    except (json.JSONDecodeError, TypeError):
                        pass

            # Update access time
            self._index.touch(block_hash)
            self._stats["loads"] += 1
            self._stats["hits"] += 1

            # Promote to hot cache for faster access next time
            if self._hot_cache_enabled and promote_to_hot_cache:
                self._promote_to_hot_cache(
                    block_hash, arrays, file_metadata, block_metadata
                )

            logger.debug(
                f"Loaded block with metadata from SSD cache: {block_hash.hex()[:16]}..."
            )
            return cache_data, metadata_dict

        except Exception as e:
            logger.error(f"Failed to load block from SSD cache: {e}")
            self._stats["errors"] += 1
            # Remove corrupted entry
            self._index.remove(block_hash)
            try:
                file_path.unlink()
            except Exception:
                pass
            return None, None

    def get_block_metadata(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
        """
        Get metadata for a block without loading the data.

        Args:
            block_hash: Content hash for the block.

        Returns:
            PagedSSDBlockMetadata if found, None otherwise.
        """
        return self._index.get(block_hash)

    def has_block(self, block_hash: bytes) -> bool:
        """
        Check if a block exists in cache (hot cache, pending writes, or SSD storage).

        Args:
            block_hash: Content hash for the block.

        Returns:
            True if block exists in hot cache, pending write buffer, or SSD index.
        """
        if self._index.contains(block_hash):
            return True
        # Block may have been evicted from SSD index but still in hot cache
        with self._hot_cache_lock:
            if block_hash in self._hot_cache:
                return True
        # Block may be evicted from hot cache and awaiting SSD write
        with self._pending_write_hashes_lock:
            if block_hash in self._pending_write_buffers:
                return True
        return False

    def preload_matched_blocks(self, block_hashes: list[bytes]) -> int:
        """
        Parallel-load matched blocks from SSD into hot cache.

        For cold-start optimization: loads blocks that exist on SSD but not
        in hot cache, using parallel I/O. After preload, subsequent
        load_block() / load_block_with_metadata() calls hit hot cache (~0ms)
        instead of SSD (~2ms per block).

        Individual block failures are non-fatal (logged and skipped).

        Args:
            block_hashes: Block hashes confirmed as cache hits.

        Returns:
            Number of blocks successfully loaded into hot cache.
        """
        if not self._hot_cache_enabled:
            return 0

        if not HAS_MLX:
            return 0

        # Filter to blocks that need loading: in SSD index but not hot cache
        to_load = []
        for bh in block_hashes:
            metadata = self._index.get(bh)
            if metadata is None:
                continue
            if self._hot_cache_get(bh) is not None:
                continue
            to_load.append((bh, metadata))

        if len(to_load) < 4:
            return 0

        # Guard: don't preload more than available hot cache capacity.
        # If we preload N blocks but hot cache can only hold M < N,
        # blocks evict each other and reconstruct_cache falls back to SSD.
        # CPD-accepted (GLM L1).
        available = self._hot_cache_available_bytes()
        if available <= 0:
            return 0
        capped_to_load: list[tuple[bytes, PagedSSDBlockMetadata]] = []
        selected_bytes = 0
        for bh, metadata in to_load:
            try:
                block_bytes = max(0, int(getattr(metadata, "file_size", 0) or 0))
            except (TypeError, ValueError):
                block_bytes = 0
            if block_bytes <= 0:
                try:
                    block_bytes = metadata.file_path.stat().st_size
                except OSError:
                    block_bytes = 0
            if selected_bytes + block_bytes > available:
                break
            capped_to_load.append((bh, metadata))
            selected_bytes += block_bytes
        to_load = capped_to_load
        if len(to_load) < 4:
            return 0

        # Cap workers to limit peak memory (each load allocates ~122-275MB).
        # 8 workers ≈ 1.4GB peak, vs 2.8GB at 16. CPD-accepted (G1/Q3).
        start = time.perf_counter()
        loaded_count = 0
        max_workers = min(8, len(to_load))

        def _load_one(block_hash: bytes, metadata: PagedSSDBlockMetadata) -> bool:
            file_path = metadata.file_path
            if not file_path.exists():
                return False
            try:
                arrays, file_metadata = mx.load(str(file_path), return_metadata=True)
                if (
                    file_metadata
                    and file_metadata.get("omlx_cache_format_version")
                    not in _READABLE_CACHE_FORMAT_VERSIONS
                ):
                    return False
                self._promote_to_hot_cache(block_hash, arrays, file_metadata, metadata)
                return True
            except Exception as e:
                logger.warning(f"Preload failed for block {block_hash.hex()[:16]}: {e}")
                return False

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(_load_one, bh, meta): bh for bh, meta in to_load}
            for future in as_completed(futures):
                try:
                    if future.result():
                        loaded_count += 1
                except Exception:
                    pass

        elapsed_ms = (time.perf_counter() - start) * 1000
        self._stats["preload_calls"] += 1
        self._stats["preload_blocks_loaded"] += loaded_count
        self._stats["preload_time_ms"] += elapsed_ms

        if loaded_count > 0:
            logger.info(
                f"Preloaded {loaded_count}/{len(to_load)} blocks into hot cache "
                f"(workers={max_workers}, time={elapsed_ms:.1f}ms)"
            )
        return loaded_count

    def adopt_layer_signature_if_unset(
        self, layer_cache_types: list[str] | None
    ) -> bool:
        """Adopt ``layer_cache_types`` as the expected signature if none was set.

        The scheduler may not be able to derive the post-patch cache layout
        before constructing this manager (TurboQuant / MTP / dtype changes
        happen at model-load time). Save sites pass the live signature on
        every call, so the manager can adopt it the first time it sees one.

        Returns True when adoption actually happened (caller can use this
        to trigger the one-shot sweep). Returns False when the manager
        already had a signature or when ``layer_cache_types`` is empty.
        """
        if not layer_cache_types:
            return False
        if self._expected_layer_cache_types is not None:
            return False
        canonical = _canonicalize_layer_cache_types(layer_cache_types)
        with self._lock:
            if self._expected_layer_cache_types is not None:
                return False  # raced with another adopter
            self._expected_layer_cache_types = list(layer_cache_types)
            self._signature_sweep_completed = False
        logger.info(
            "PagedSSDCacheManager adopted layer cache signature "
            "(%d layers, %d unique types)",
            len(layer_cache_types),
            len(set(canonical or ())),
        )
        return True

    def set_expected_layer_signature(
        self,
        layer_cache_types: list[str] | None,
        *,
        turboquant_kv_bits: float | None = None,
        cachelist_subtypes: dict[str, list[str]] | None = None,
    ) -> bool:
        """Set the live layer-cache signature, replacing stale expectations.

        Unlike ``adopt_layer_signature_if_unset``, this is used by callers that
        learn the final cache layout after manager construction (for example
        TurboQuant settings applied by the engine after the scheduler starts).

        ``turboquant_kv_bits`` is the live TurboQuant bit depth (None when
        TurboQuant is inactive). A bit-depth change alone also triggers the
        sweep: blocks written at another depth have a different packed state
        width and would crash batch concatenation if mixed (#2045).

        Returns True when the canonical signature changed and a stale-signature
        sweep should run. Returns False for empty input or a canonical no-op.
        """
        if not layer_cache_types:
            return False

        new_signature = list(layer_cache_types)
        new_canonical = _canonicalize_layer_cache_types(new_signature)
        new_bits = (
            float(turboquant_kv_bits) if turboquant_kv_bits is not None else None
        )

        with self._lock:
            old_signature = self._expected_layer_cache_types
            old_canonical = _canonicalize_layer_cache_types(old_signature)
            bits_changed = new_bits != self._expected_turboquant_kv_bits
            subtypes_changed = (
                cachelist_subtypes != self._expected_cachelist_subtypes
            )
            if (
                old_canonical == new_canonical
                and not bits_changed
                and not subtypes_changed
            ):
                if old_signature != new_signature:
                    self._expected_layer_cache_types = new_signature
                return False

            self._expected_layer_cache_types = new_signature
            self._expected_turboquant_kv_bits = new_bits
            self._expected_cachelist_subtypes = cachelist_subtypes
            self._signature_sweep_completed = False

        logger.info(
            "PagedSSDCacheManager updated layer cache signature "
            "(%d layers, %d unique types, turboquant_kv_bits=%s, "
            "cachelist_subtypes=%s)",
            len(new_signature),
            len(set(new_canonical or ())),
            new_bits,
            "yes" if cachelist_subtypes else "no",
        )
        return True

    def invalidate_stale_layer_signature(self) -> int:
        """Drop in-memory index entries whose layer_cache_types — or, when a
        TurboQuant depth is expected, whose recorded bit depth — disagree
        with the current expected signature.

        Scoped to the current ``_expected_model_name``: blocks belonging to
        other models share the SSD directory and remain valid for them, so
        we leave them alone. Legacy blocks without a recorded ``model_name``
        are also skipped — we cannot safely attribute them.

        The SSD files are unlinked from the index (and any in-memory hot
        copy is dropped), but the on-disk file is left for LRU to reclaim
        later. Returns the number of blocks dropped from the index.

        Idempotent: a second call after a clean sweep returns 0.
        """
        if self._expected_layer_cache_types is None:
            return 0
        if not self._expected_model_name:
            # Without an owning model_name we cannot scope safely; refuse
            # rather than risk evicting another model's blocks.
            return 0
        if self._signature_sweep_completed:
            return 0

        expected = _canonicalize_layer_cache_types(self._expected_layer_cache_types)
        expects_bits = self._expected_turboquant_kv_bits is not None

        with self._index._lock:
            stale: list[bytes] = []
            for h, meta in self._index._index.items():
                if not meta.model_name or meta.model_name != self._expected_model_name:
                    continue
                got = _canonicalize_layer_cache_types(meta.layer_cache_types)
                if got is None:
                    # Pre-signature blocks lack the metadata to judge the
                    # layout. Without a depth expectation, skip rather than
                    # guess — newer saves will replace them. With one, the
                    # block can no more prove its packed width than its
                    # layout, so it is unsafe to keep (see
                    # _signature_bits_match).
                    if expects_bits:
                        stale.append(h)
                    continue
                if got != expected:
                    stale.append(h)
                    continue
                if not self._signature_bits_match(meta.cache_signature):
                    stale.append(h)
                    continue
                if self._expected_cachelist_subtypes is not None and (
                    _signature_cachelist_subtypes(meta.cache_signature)
                    != self._expected_cachelist_subtypes
                ):
                    stale.append(h)

        for h in stale:
            self.forget_block(h)

        self._signature_sweep_completed = True

        if stale:
            logger.info(
                "Invalidated %d SSD index entries with stale layer "
                "cache signature for model %r (kept %d)",
                len(stale),
                self._expected_model_name,
                len(self._index._index),
            )
        return len(stale)

    def forget_block(self, block_hash: bytes) -> bool:
        """
        Remove a block from this manager's in-memory indexes without deleting
        its SSD file.

        Used when a prefix entry points at a block that is incompatible with
        the current model/layout. The file may still be valid for another
        model sharing the same cache directory.
        """
        with self._lock:
            removed = self._hot_cache_remove(block_hash) is not None

            with self._pending_write_hashes_lock:
                if block_hash in self._pending_write_buffers:
                    removed = True
                self._pending_write_buffers.pop(block_hash, None)
                self._pending_write_hashes.discard(block_hash)

            metadata = self._index.remove(block_hash)
            if metadata is not None:
                self._incompatible_index.add(metadata)
                removed = True

            return removed

    def delete_block(self, block_hash: bytes) -> bool:
        """
        Delete a block from SSD storage.

        Args:
            block_hash: Content hash for the block.

        Returns:
            True if deleted successfully.
        """
        with self._lock:
            # Also remove from hot cache
            self._hot_cache_remove(block_hash)

            # Also remove from pending write buffer
            with self._pending_write_hashes_lock:
                self._pending_write_buffers.pop(block_hash, None)
                self._pending_write_hashes.discard(block_hash)

            metadata = self._index.remove(block_hash)
            incompatible_metadata = self._incompatible_index.remove(block_hash)
            if metadata is None:
                metadata = incompatible_metadata
            if metadata is None:
                return False

            try:
                if metadata.file_path.exists():
                    metadata.file_path.unlink()
                    logger.debug(f"Deleted SSD cache file: {metadata.file_path}")
                return True
            except Exception as e:
                logger.error(f"Failed to delete SSD cache file: {e}")
                return False

    # Use at most 99% of available disk space to avoid filling disk completely
    _DISK_SAFE_RATIO = 0.99

    def _get_effective_max_size(self) -> int:
        """Get effective max size considering actual disk free space.

        Returns the minimum of configured max_size and 99% of disk space
        available for cache (current cache size + disk free). This ensures
        eviction triggers before the disk fills up even when other processes
        consume disk space after the server started.

        Uses a 30-second TTL cache for shutil.disk_usage() results.
        """
        if self._cache_dir is None:
            return self._max_size

        # Take the lock so a concurrent writer-thread invalidation
        # (sets _disk_usage_cache=None on ENOSPC) can't interleave with
        # this read-check-write and let one save see a fresh value paired
        # with a stale timestamp (or vice versa).
        now = time.monotonic()
        with self._lock:
            if (
                self._disk_usage_cache is None
                or now - self._disk_usage_cache_time > 30.0
            ):
                try:
                    self._disk_usage_cache = shutil.disk_usage(self._cache_dir)
                except OSError as e:
                    logger.warning(
                        f"Failed to check disk usage for SSD cache dir "
                        f"{self._cache_dir}: {e}"
                    )
                    return self._max_size
                self._disk_usage_cache_time = now
            disk_free = self._disk_usage_cache.free

        disk_available = self._tracked_ssd_size() + disk_free
        disk_limit = int(disk_available * self._DISK_SAFE_RATIO)
        return min(self._max_size, disk_limit)

    def _evict_tracked_until_size(
        self,
        target_size: int,
        max_count: int | None = None,
    ) -> list[tuple[PagedSSDCacheIndex, PagedSSDBlockMetadata]]:
        """Remove oldest tracked SSD entries from their indexes until target."""
        evicted: list[tuple[PagedSSDCacheIndex, PagedSSDBlockMetadata]] = []

        while self._tracked_ssd_size() > target_size:
            if max_count is not None and len(evicted) >= max_count:
                break

            compatible = self._index.get_lru_entries(1)
            incompatible = self._incompatible_index.get_lru_entries(1)
            if not compatible and not incompatible:
                break

            if compatible and incompatible:
                if incompatible[0].last_access <= compatible[0].last_access:
                    source_index = self._incompatible_index
                    candidate = incompatible[0]
                else:
                    source_index = self._index
                    candidate = compatible[0]
            elif incompatible:
                source_index = self._incompatible_index
                candidate = incompatible[0]
            else:
                source_index = self._index
                candidate = compatible[0]

            metadata = source_index.remove(candidate.block_hash)
            if metadata is not None:
                evicted.append((source_index, metadata))

        return evicted

    def _enforce_size_limit_for_new_block(
        self,
        estimated_new_size: int = 1 * 1024 * 1024,
        *,
        max_unlinks: int | None = None,
        unbounded: bool = False,
    ) -> None:
        """Enforce size limit before adding a new block.

        ``estimated_new_size`` should be the actual byte size of the block
        about to be inserted. The 1 MiB default is for callers that don't
        yet know the size at the time eviction is needed; passing the
        actual size avoids cache oscillation around the configured limit.
        """
        effective_max = self._get_effective_max_size()

        # Warn when disk pressure shrinks effective limit well below configured
        # (throttled to once per 60s to avoid log spam)
        if effective_max < self._max_size * 0.1:
            now = time.monotonic()
            if now - self._last_disk_pressure_warn > 60.0:
                self._last_disk_pressure_warn = now
                logger.warning(
                    f"SSD cache disk pressure: effective limit "
                    f"{format_bytes(effective_max)} "
                    f"(configured {format_bytes(self._max_size)}), "
                    f"disk nearly full"
                )
        target_size = effective_max - estimated_new_size
        if target_size < 0:
            target_size = int(effective_max * 0.9)

        max_count = None if unbounded else max_unlinks
        if max_count is None and not unbounded:
            max_count = _MAX_INLINE_UNLINKS_PER_SAVE

        if self._tracked_ssd_size() > target_size:
            evicted = self._evict_tracked_until_size(
                target_size,
                max_count=max_count,
            )
            # Inline unlinks on the calling thread. Eviction typically returns
            # a single entry per save because the tracked LRU walk stops as
            # soon as the shared SSD budget is back under target. Inline
            # removes bounded-queue contention entirely. Hot cache is NOT
            # touched here — ``delete_block()`` is the only path that clears
            # both tiers.
            #
            # Bounded inline burst. The ENOSPC-recovery path invalidates the
            # 30 s disk-usage cache, which can shrink the next
            # ``_get_effective_max_size`` call sharply. Cap the burst at
            # ``_MAX_INLINE_UNLINKS_PER_SAVE`` and leave remaining LRU
            # entries in their indexes so subsequent saves drain the rest.
            # Bounds per-call latency at the cost of taking multiple saves
            # to fully reconverge.
            for source_index, metadata in evicted:
                self._unlink_evicted(metadata, source_index)
            if max_count is not None and len(evicted) >= max_count:
                logger.debug(
                    f"Inline eviction capped at {max_count} entries; "
                    f"{self._tracked_ssd_size() - target_size} bytes remain "
                    f"above target for subsequent saves to drain"
                )

    def enforce_size_limit(self) -> int:
        """
        Enforce SSD cache size limit by evicting LRU files.

        Returns:
            Number of bytes freed.
        """
        # Decide what to evict under the lock, but perform unlinks outside
        # it: a single unlink on a slow disk (NFS / encrypted FS / ENOSPC
        # retry path) can block tens to hundreds of ms, and every
        # _get_effective_max_size() / writer-thread cache-invalidation
        # contends on self._lock. The index has its own internal lock
        # protecting the LRU/size accounting.
        with self._lock:
            initial_size = self._tracked_ssd_size()
            effective_max = self._get_effective_max_size()

            if initial_size <= effective_max:
                return 0

            target_size = int(effective_max * 0.9)  # 90% of effective max
            evicted = self._evict_tracked_until_size(target_size)

        # Do NOT remove from hot cache — see _enforce_size_limit_for_new_block
        for source_index, metadata in evicted:
            self._unlink_evicted(metadata, source_index)

        freed = initial_size - self._tracked_ssd_size()
        logger.info(
            f"SSD cache size enforcement: freed {format_bytes(freed)}, "
            f"evicted {len(evicted)} files"
        )
        return freed

    def _unlink_evicted(
        self,
        metadata: PagedSSDBlockMetadata,
        source_index: PagedSSDCacheIndex | None = None,
    ) -> None:
        """Delete an evicted block file from disk.

        On unlink failure other than FileNotFoundError, re-add the
        metadata to the index so ``total_size`` keeps reflecting actual
        on-disk bytes; without this, accumulated failures would let the
        cache silently exceed ``max_size`` (the index would report free
        space that does not exist on disk).
        """
        try:
            metadata.file_path.unlink(missing_ok=True)
            self._stats["evictions"] += 1
        except OSError as e:
            restore_index = source_index or self._index
            # Restore the index entry so total_size matches disk reality.
            # The re-added entry lands at the LRU tail (most-recently
            # touched), which deprioritises immediate re-eviction.
            restore_index.add(metadata)
            self._stats["evict_unlink_failures"] += 1
            logger.exception(
                "Failed to delete evicted SSD cache file %s: %s",
                metadata.file_path,
                e,
            )

    def clear_hot_cache(self) -> int:
        """Clear all in-memory (hot) cache entries.

        Returns:
            Number of entries cleared.
        """
        with self._hot_cache_lock:
            count = len(self._hot_cache)
            self._hot_cache.clear()
            self._hot_cache_total_bytes = 0
        if self._hot_cache_budget is not None:
            self._hot_cache_budget.forget_owner(self)
        if count:
            logger.info("Cleared %d hot cache entries", count)
        return count

    def shrink_hot_cache_to(
        self,
        target_bytes: int,
        protected_hashes: set[bytes] | None = None,
    ) -> int:
        """Shrink this manager's hot cache to ``target_bytes`` by local LRU."""
        target_bytes = max(0, int(target_bytes))
        protected_hashes = protected_hashes or set()

        if self._hot_cache_budget is not None:
            return self._hot_cache_budget.shrink_to(
                target_bytes, protected_hashes=protected_hashes
            )

        evicted_entries: list[tuple[bytes, dict, int]] = []
        with self._hot_cache_lock:
            while self._hot_cache_total_bytes > target_bytes and self._hot_cache:
                victim_hash = None
                for block_hash in self._hot_cache:
                    if block_hash not in protected_hashes:
                        victim_hash = block_hash
                        break
                if victim_hash is None:
                    break

                evicted = self._hot_cache.pop(victim_hash)
                size = self._hot_cache_entry_size(evicted)
                self._hot_cache_total_bytes = max(0, self._hot_cache_total_bytes - size)
                evicted_entries.append((victim_hash, evicted, size))

        freed = 0
        for block_hash, evicted, size in evicted_entries:
            freed += size
            self._handle_hot_cache_eviction(block_hash, evicted)

        if freed and self._hot_cache_only:
            logger.warning(
                "Shrank hot-cache-only tier by %s; evicted chains are not "
                "persisted to SSD",
                format_bytes(freed),
            )
        elif freed:
            logger.info("Shrank hot cache by %s", format_bytes(freed))
        return freed

    def clear(self) -> int:
        """
        Clear all SSD cache files.

        Returns:
            Number of files deleted.
        """
        with self._lock:
            count = 0
            block_hashes = (
                self._index.get_all_hashes() + self._incompatible_index.get_all_hashes()
            )
            for block_hash in dict.fromkeys(block_hashes):
                if self.delete_block(block_hash):
                    count += 1

            logger.info(f"Cleared SSD cache: deleted {count} files")
            return count

    def get_stats(self) -> PagedSSDCacheStats:
        """
        Get SSD cache statistics.

        Returns:
            PagedSSDCacheStats with cache metrics.
        """
        with self._lock:
            with self._hot_cache_lock:
                hot_entries = len(self._hot_cache)
                hot_size = self._hot_cache_total_bytes
            return PagedSSDCacheStats(
                hits=self._stats["hits"],
                misses=self._stats["misses"],
                evictions=self._stats["evictions"],
                saves=self._stats["saves"],
                saves_persisted=self._stats["saves_persisted"],
                loads=self._stats["loads"],
                errors=self._stats["errors"],
                evict_unlink_failures=self._stats["evict_unlink_failures"],
                total_size_bytes=self._tracked_ssd_size(),
                max_size_bytes=self._get_effective_max_size(),
                configured_max_size_bytes=self._max_size,
                num_files=self._tracked_ssd_count(),
                hot_cache_entries=hot_entries,
                hot_cache_size_bytes=hot_size,
                hot_cache_max_bytes=self._effective_hot_cache_max_bytes(),
                hot_cache_hits=self._stats["hot_cache_hits"],
                hot_cache_evictions=self._stats["hot_cache_evictions"],
                hot_cache_promotions=self._stats["hot_cache_promotions"],
                ssd_write_drops=self._stats["ssd_write_drops"],
                ssd_inline_write_fallbacks=self._stats["ssd_inline_write_fallbacks"],
            )

    def get_stats_for_model(self, model_name: str) -> PagedSSDCacheStats:
        """Get model-scoped SSD cache statistics.

        The SSD cache directory can be shared across multiple loaded models, so
        dashboard per-model rows must be filtered by block metadata rather than
        reusing the global cache totals.
        """
        normalized_name = model_name.rstrip("/")
        basename = os.path.basename(normalized_name) if normalized_name else ""

        def _matches(candidate: str) -> bool:
            candidate = candidate.rstrip("/")
            if not candidate:
                return False
            if candidate == normalized_name:
                return True
            if basename and os.path.basename(candidate) == basename:
                return True
            return False

        with self._lock:
            indexed_entries = [
                metadata
                for metadata in self._index.get_all_metadata()
                if _matches(metadata.model_name)
            ]
            indexed_size = sum(metadata.file_size for metadata in indexed_entries)
            indexed_count = len(indexed_entries)

            with self._hot_cache_lock:
                hot_entries = []
                hot_size = 0
                for entry in self._hot_cache.values():
                    blk_meta = entry.get("block_metadata")
                    if blk_meta is None or not _matches(blk_meta.model_name):
                        continue
                    hot_entries.append(entry)
                    hot_size += self._hot_cache_entry_size(entry)

            return PagedSSDCacheStats(
                hits=self._stats["hits"],
                misses=self._stats["misses"],
                evictions=self._stats["evictions"],
                saves=self._stats["saves"],
                saves_persisted=self._stats["saves_persisted"],
                loads=self._stats["loads"],
                errors=self._stats["errors"],
                evict_unlink_failures=self._stats["evict_unlink_failures"],
                total_size_bytes=indexed_size,
                max_size_bytes=self._get_effective_max_size(),
                configured_max_size_bytes=self._max_size,
                num_files=indexed_count,
                hot_cache_entries=len(hot_entries),
                hot_cache_size_bytes=hot_size,
                hot_cache_max_bytes=self._effective_hot_cache_max_bytes(),
                hot_cache_hits=self._stats["hot_cache_hits"],
                hot_cache_evictions=self._stats["hot_cache_evictions"],
                hot_cache_promotions=self._stats["hot_cache_promotions"],
                ssd_write_drops=self._stats["ssd_write_drops"],
                ssd_inline_write_fallbacks=self._stats["ssd_inline_write_fallbacks"],
            )

    def get_stats_dict(self) -> dict[str, Any]:
        """
        Get SSD cache statistics as a dictionary.

        This method provides the legacy dictionary format for compatibility.

        Returns:
            Dictionary with cache statistics.
        """
        with self._lock:
            with self._hot_cache_lock:
                hot_entries = len(self._hot_cache)
                hot_size = self._hot_cache_total_bytes
            effective_max = self._get_effective_max_size()
            return {
                "cache_dir": str(self._cache_dir) if self._cache_dir else "None",
                "max_size": effective_max,
                "max_size_formatted": format_bytes(effective_max),
                "configured_max_size": self._max_size,
                "configured_max_size_formatted": format_bytes(self._max_size),
                "total_size": self._tracked_ssd_size(),
                "total_size_formatted": format_bytes(self._tracked_ssd_size()),
                "utilization": (
                    self._tracked_ssd_size() / effective_max
                    if effective_max > 0
                    else 0.0
                ),
                "num_files": self._tracked_ssd_count(),
                "hot_cache_entries": hot_entries,
                "hot_cache_size_bytes": hot_size,
                "hot_cache_max_bytes": self._effective_hot_cache_max_bytes(),
                "hot_cache_size_formatted": format_bytes(hot_size),
                "hot_cache_max_formatted": format_bytes(
                    self._effective_hot_cache_max_bytes()
                ),
                **self._stats,
            }

    def close(self) -> None:
        """Close the SSD cache manager, flushing hot cache and pending writes."""
        logger.info("Shutting down PagedSSDCacheManager...")

        # Flush hot cache entries to SSD before shutdown.
        # Dirty blocks wait for queue space first; sustained saturation falls
        # back to an inline write on this thread.
        if self._hot_cache_enabled:
            with self._hot_cache_lock:
                entries_to_flush = list(self._hot_cache.items())
            flushed = 0
            failed = 0
            for block_hash, entry in entries_to_flush:
                if self._writer_thread and not self._writer_thread.is_alive():
                    logger.warning(
                        "Writer thread died during shutdown flush, "
                        f"aborting ({flushed} flushed, "
                        f"{len(entries_to_flush) - flushed - failed} remaining)"
                    )
                    break
                blk_meta = entry.get("block_metadata")
                if not entry.get("dirty", True):
                    continue
                if blk_meta and blk_meta.file_path.exists():
                    continue
                if self._enqueue_ssd_write(block_hash, entry, blocking=True):
                    flushed += 1
                else:
                    failed += 1
            if flushed:
                logger.info(f"Flushed {flushed} hot cache blocks to SSD")
            if failed:
                logger.warning(f"Failed to flush {failed} hot cache blocks")

        # Signal writer thread to stop (after processing remaining queue)
        if self._writer_thread:
            self._writer_shutdown.set()

            # Send sentinel to unblock the writer if it's waiting on the queue
            try:
                self._write_queue.put_nowait(None)
            except queue.Full:
                pass  # Writer will check shutdown flag on next iteration

            # Wait for writer to finish — longer timeout to allow flush
            timeout = 120 if self._hot_cache_enabled else 60
            self._writer_thread.join(timeout=timeout)
            if self._writer_thread.is_alive():
                logger.warning(
                    f"SSD cache writer thread did not stop within {timeout}s"
                )

        # Clear hot cache and pending write buffer
        with self._hot_cache_lock:
            self._hot_cache.clear()
            self._hot_cache_total_bytes = 0
        if self._hot_cache_budget is not None:
            self._hot_cache_budget.forget_owner(self)
        with self._pending_write_hashes_lock:
            self._pending_write_buffers.clear()
            self._pending_write_hashes.clear()

        logger.debug("PagedSSDCacheManager closed")

    def __repr__(self) -> str:
        return (
            f"PagedSSDCacheManager(dir={self._cache_dir}, "
            f"size={format_bytes(self._tracked_ssd_size())}/"
            f"{format_bytes(self._max_size)}, "
            f"files={self._tracked_ssd_count()})"
        )

    # =========================================================================
    # CacheManager ABC Interface Implementation
    # =========================================================================

    def fetch(self, key: Any) -> tuple[Any | None, bool]:
        """
        Fetch a cached block from SSD storage.

        Args:
            key: Block hash (bytes) to look up.

        Returns:
            Tuple of (cache_data, True) if found, (None, False) otherwise.
        """
        if not isinstance(key, bytes):
            return None, False

        cache_data = self.load_block(key)
        if cache_data is not None:
            return cache_data, True
        return None, False

    def store(self, key: Any, value: Any) -> bool:
        """
        Store a block in SSD cache.

        Args:
            key: Block hash (bytes).
            value: Tuple of (cache_data, token_count) or just cache_data.

        Returns:
            True if stored successfully.
        """
        if not isinstance(key, bytes):
            return False

        if isinstance(value, tuple) and len(value) >= 2:
            cache_data, token_count = value[0], value[1]
            model_name = value[2] if len(value) > 2 else ""
        else:
            cache_data = value
            token_count = 0
            model_name = ""

        return self.save_block(key, cache_data, token_count, model_name)

    def evict(self, key: Any) -> bool:
        """
        Evict a specific block from SSD cache.

        Args:
            key: Block hash (bytes) to evict.

        Returns:
            True if evicted, False if not found.
        """
        if not isinstance(key, bytes):
            return False

        return self.delete_block(key)

    @property
    def size(self) -> int:
        """
        Get the current number of cached blocks.

        Returns:
            Number of cached blocks.
        """
        return self._index.count

    @property
    def max_size(self) -> int:
        """
        Get the effective maximum cache size in bytes.

        This accounts for actual disk free space, returning the minimum of
        the configured max size and 99% of available disk space for cache.

        Returns:
            Effective maximum cache size in bytes.
        """
        return self._get_effective_max_size()

    @property
    def configured_max_size(self) -> int:
        """
        Get the originally configured maximum cache size in bytes.

        Returns:
            Configured maximum cache size in bytes.
        """
        return self._max_size


# ===== FILE: omlx/cache/pooling_delta.py =====

# SPDX-License-Identifier: Apache-2.0
"""Storage-only delta encoding for cumulative PoolingCache snapshots."""

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)

POOLING_CACHE_DELTA_CLASS = "PoolingCacheDelta"
POOLING_CACHE_DELTA_FORMAT_VERSION = "4"


def compact_pooling_cache_snapshot(
    extracted: list[dict[str, Any]],
    token_count: int,
    block_size: int,
) -> list[dict[str, Any]]:
    """Replace cumulative PoolingCache tensors with the current block delta.

    PoolingCache is non-sliceable because its buffer and gate are required for
    continued decoding, but its ``pooled`` tensor is append-only. Boundary
    snapshots therefore only need the newly appended pooled rows. The absolute
    row range is recorded so restore can reject incomplete or out-of-order
    delta chains instead of silently constructing a corrupt cache.

    The input is modified in place and returned for convenience. If an
    unfamiliar state shape or length is encountered, that sub-cache remains a
    legacy full snapshot.
    """
    if token_count <= 0 or block_size <= 0:
        return extracted

    for layer in extracted:
        if layer.get("class_name") != "CacheList":
            continue

        sub_states = layer.get("state")
        class_names = layer.get("sub_class_names")
        meta_state = layer.get("meta_state")
        if (
            not class_names
            and isinstance(meta_state, (list, tuple))
            and len(meta_state) >= 1
        ):
            class_names = meta_state[0]
        sub_meta_states = (
            meta_state[1]
            if isinstance(meta_state, (list, tuple)) and len(meta_state) >= 2
            else []
        )
        if not isinstance(sub_states, (list, tuple)) or not isinstance(
            class_names, (list, tuple)
        ):
            continue

        delta_ranges: dict[str, list[int]] = {}
        compacted_states = list(sub_states)
        for sub_idx, (class_name, state) in enumerate(
            zip(class_names, sub_states, strict=False)
        ):
            if class_name != "PoolingCache" or not isinstance(state, (list, tuple)):
                continue
            if len(state) not in (3, 5) or sub_idx >= len(sub_meta_states):
                continue

            try:
                ratio = int(sub_meta_states[sub_idx])
                if ratio <= 0:
                    continue
                pooled = state[2]
                pooled_length = int(pooled.shape[1])
                expected_end = token_count // ratio
                expected_start = max(0, token_count - block_size) // ratio
            except (AttributeError, TypeError, ValueError, IndexError):
                continue

            if pooled_length != expected_end:
                logger.debug(
                    "Keeping full PoolingCache snapshot at layer %s sub-cache "
                    "%d: pooled=%s expected=%s ratio=%s",
                    layer.get("layer_idx"),
                    sub_idx,
                    pooled_length,
                    expected_end,
                    ratio,
                )
                continue

            delta = pooled[:, expected_start:expected_end]
            compacted_states[sub_idx] = (
                state[0],
                state[1],
                delta,
                *state[3:],
            )
            delta_ranges[str(sub_idx)] = [expected_start, expected_end]

        if delta_ranges:
            layer["state"] = compacted_states
            layer["pooling_delta_ranges"] = delta_ranges

    return extracted


# ===== FILE: omlx/cache/prefix_cache.py =====

# SPDX-License-Identifier: Apache-2.0
"""
Block-Aware Prefix Cache for oMLX.

Provides prefix caching using PagedCacheManager for block-based storage
with SSD persistence. oMLX only supports paged SSD-based caching.
"""

import logging
import math
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

try:
    import mlx.core as mx

    HAS_MLX = True
except ImportError:
    HAS_MLX = False

from ._rotating_subclass import PrefillReadyRotatingKVCache
from .hybrid_cache import ModelCacheConfig
from .interface import CacheManager
from .paged_cache import (
    BlockTable,
    CacheBlock,
    PagedCacheManager,
    compute_block_hash,
    resolve_block_extra_keys,
)
from .paged_ssd_cache import PagedSSDCacheManager
from .pooling_delta import POOLING_CACHE_DELTA_CLASS
from .stats import PrefixCacheStats
from .type_registry import CacheTypeRegistry

logger = logging.getLogger(__name__)

# Cap on the supersede-on-extend lineage map (tip hash -> previous tip hash).
# Each entry is two 32-byte hashes; the cap only guards against unbounded
# growth from many distinct conversation chains over a long-lived process.
_TIP_LINEAGE_MAX_ENTRIES = 4096


@dataclass
class BlockCacheEntry:
    """Entry mapping a token sequence to cache blocks."""

    block_table: BlockTable
    last_access: float


class BlockAwarePrefixCache(CacheManager):
    """
    Prefix cache that uses PagedCacheManager for block-based storage.

    Features:
    - Block-level prefix sharing (256 tokens per block)
    - paged SSD-only storage via PagedSSDCacheManager
    - Hash-based deduplication across requests
    - Reference counting for memory efficiency

    Implements the CacheManager ABC interface for consistency with other
    cache implementations in oMLX.

    In paged SSD-only mode:
    - All KV cache data is stored on paged SSD via PagedSSDCacheManager
    - PagedCacheManager only stores metadata (no cache_data in blocks)
    - Cache data is loaded from paged SSD when needed for inference

    Example:
        cold_manager = PagedSSDCacheManager(cache_dir=Path("/tmp/cache"), ...)
        paged_manager = PagedCacheManager(block_size=256, max_blocks=1000)
        cache = BlockAwarePrefixCache(model, paged_manager, cold_manager)

        # Check for cached prefix
        block_table, remaining_tokens = cache.fetch_cache(request_id, tokens)

        # After generation, store cache
        cache.store_cache(request_id, tokens, kv_cache_data)

        # Clean up when request completes
        cache.release_cache(request_id)
    """

    def __init__(
        self,
        model: Any,
        paged_cache_manager: PagedCacheManager,
        paged_ssd_cache_manager: PagedSSDCacheManager | None = None,
    ):
        """
        Initialize block-aware prefix cache.

        Args:
            model: The MLX model (used for identification)
            paged_cache_manager: The PagedCacheManager instance for block management
            paged_ssd_cache_manager: The PagedSSDCacheManager for SSD storage (required for paged SSD-only mode)
        """
        self.model = model
        self.model_key = id(model)
        self.paged_cache = paged_cache_manager
        self.paged_ssd_cache = paged_ssd_cache_manager
        self.block_size = paged_cache_manager.block_size

        # Expected number of layers for cache validation
        self.expected_num_layers = self._get_model_num_layers(model)

        # Hash table for quick prefix lookup
        # Maps chain-hash(prefix) -> (prefix_len, block_ids, num_blocks)
        self._prefix_index: dict[bytes, tuple[int, tuple[int, ...], int]] = {}

        # Tie the index lifecycle to the paged cache's hash associations.
        # Without these hooks the index only ever grew (entries were dropped
        # solely by clear()), leaking Python heap for the process lifetime
        # on long-uptime servers.
        paged_cache_manager.on_block_hash_dropped = self._on_block_hash_dropped
        paged_cache_manager.on_hash_map_cleared = self._on_hash_map_cleared

        # Request to block table mapping
        self._request_tables: dict[str, BlockCacheEntry] = {}

        # Supersede-on-extend lineage for rotating (sliding-window) models:
        # newest tip block hash -> previous tip block hash. When a chain is
        # extended again, the entry two generations back is stripped of its
        # rotating payload (see _strip_rotating_payload); the immediate
        # previous tip is kept intact as the walk-back fallback.
        self._rotating_tip_lineage: dict[bytes, bytes] = {}

        # Callback for restoring cold blocks (deprecated in paged SSD-only mode)
        # Kept for API compatibility
        self._cold_restore_callback: Callable[[int, bytes], bool] | None = None

        # Statistics
        self._hits = 0
        self._misses = 0
        self._tokens_saved = 0
        self._partial_block_skips = 0
        self._partial_tokens_skipped = 0
        self._tokens_matched_total = 0
        self._tokens_requested_total = 0
        self._last_partial_tokens_skipped = 0
        self._last_tokens_to_next_block = 0

    def _get_model_num_layers(self, model: Any) -> int:
        """
        Get the expected number of *cache layers* for validation.

        For hybrid models, the number of cache entries (from ``make_cache()``)
        may be smaller than the architectural layer count (``model.layers``),
        because some layer types do not produce cache state.

        Args:
            model: The MLX model

        Returns:
            Number of cache layers, or 0 if cannot be determined
        """
        # Prefer cache-layer count when available (hybrid-model safe).
        make_cache = getattr(model, "make_cache", None)
        if callable(make_cache):
            try:
                cache_list = make_cache()
                if isinstance(cache_list, list) and len(cache_list) > 0:
                    return len(cache_list)
            except Exception as e:
                logger.debug(
                    f"Could not determine cache layer count via make_cache(): {e}"
                )

        # Fallback to architectural layer count for non-hybrid models.
        if hasattr(model, "layers"):
            return len(model.layers)
        if hasattr(model, "args") and hasattr(model.args, "num_hidden_layers"):
            return model.args.num_hidden_layers
        if hasattr(model, "config") and hasattr(model.config, "num_hidden_layers"):
            return model.config.num_hidden_layers

        # Cannot determine, return 0 to skip validation
        logger.debug(
            "Cannot determine model/cache num_layers, cache layer validation disabled"
        )
        return 0

    def set_paged_ssd_cache_manager(
        self, paged_ssd_cache_manager: PagedSSDCacheManager | None
    ) -> None:
        """
        Set the PagedSSDCacheManager for SSD storage.

        This allows setting the SSD cache after initialization,
        which is useful when the scheduler creates it later.

        Args:
            paged_ssd_cache_manager: The PagedSSDCacheManager instance.
        """
        self.paged_ssd_cache = paged_ssd_cache_manager
        if paged_ssd_cache_manager is not None:
            # If the manager already knows its layer-cache signature (e.g.,
            # eager scheduler plumb-through), purge any indexed blocks left
            # over from a prior cache-config run for the current model.
            # When the signature is unset, this is a no-op; the manager
            # adopts on the first save and sweeps then.
            try:
                paged_ssd_cache_manager.invalidate_stale_layer_signature()
            except Exception as e:
                logger.warning("Stale-signature sweep on manager attach failed: %s", e)
            logger.info("PagedSSDCacheManager connected to BlockAwarePrefixCache")

    def _forget_incompatible_ssd_block(
        self,
        block_hash: bytes | None,
        block_id: int | None = None,
    ) -> None:
        """Remove an incompatible block from this manager without unlinking it.

        The SSD cache directory can be shared by multiple loaded models. A
        block that is stale for this prefix cache may still be valid for
        another model, so mismatch handling must clear local indexes only.
        """
        if block_hash is None:
            return
        try:
            if block_id is None:
                block = self.paged_cache.cached_block_hash_to_block.get_block(
                    block_hash
                )
                block_id = block.block_id if block is not None else None
            if block_id is not None:
                self.paged_cache.cached_block_hash_to_block.pop(block_hash, block_id)
                # Keep the prefix index in step with the hash map, or the
                # index keeps retrying this dead chain until the block is
                # freed (same last-block-standing rule as the paged hooks).
                self.paged_cache._notify_hash_dropped(block_hash)
        except Exception as e:
            logger.debug(f"Failed to forget incompatible paged block: {e}")
        if self.paged_ssd_cache is None:
            return
        try:
            self.paged_ssd_cache.forget_block(block_hash)
        except Exception as e:
            logger.debug(f"Failed to forget incompatible SSD block: {e}")

    @staticmethod
    def _canonical_layer_cache_types(
        layer_cache_types: list[str] | tuple[str, ...] | None,
    ) -> list[str] | None:
        """Normalize wrapper class names for metadata compatibility checks.

        Thin wrapper around the canonical implementation in
        :mod:`omlx.cache.paged_ssd_cache` so the two sites that compare
        signatures stay in lock-step.
        """
        from .paged_ssd_cache import _canonicalize_layer_cache_types

        return _canonicalize_layer_cache_types(layer_cache_types)

    def _detect_window_padding_from_blocks(
        self,
        block_ids: list[int],
    ) -> ModelCacheConfig | None:
        """Detect if blocks contain RotatingKVCache data and build config for padding.

        Checks block metadata from SSD to determine if the cached model uses
        RotatingKVCache layers. If so, builds a ModelCacheConfig with window_size
        for use with _apply_window_padding().

        Args:
            block_ids: List of block IDs to check

        Returns:
            ModelCacheConfig if RotatingKVCache detected, None otherwise
        """
        if not block_ids or self.paged_ssd_cache is None:
            return None

        first_block = self.paged_cache.allocated_blocks.get(block_ids[0])
        if not first_block or not first_block.block_hash:
            return None

        _, metadata = self.paged_ssd_cache.load_block_with_metadata(
            first_block.block_hash
        )
        if not metadata:
            return None

        layer_cache_types = metadata.get("layer_cache_types")
        # Note: CacheList layers containing RotatingKVCache sub-caches do NOT need
        # window padding. CacheList uses last-block-only storage with reject-on-partial
        # strategy, so the sliding window state is either fully restored (exact match)
        # or the entire cache is rejected (partial match).
        if not layer_cache_types or not any(
            CacheTypeRegistry.is_rotating_family(t) for t in layer_cache_types
        ):
            return None

        model_cache_config = ModelCacheConfig.from_type_list(
            layer_cache_types, model_name=""
        )

        # Extract window_size from layer meta_states
        layer_meta_states = metadata.get("layer_meta_states", [])
        max_window_size = 0
        for idx, meta in enumerate(layer_meta_states):
            if not meta or len(meta) < 2:
                continue
            # Check if this layer is RotatingKVCache
            if idx < len(layer_cache_types) and CacheTypeRegistry.is_rotating_family(
                layer_cache_types[idx]
            ):
                # RotatingKVCache meta_state: (keep, max_size, offset, _idx)
                window_size = int(meta[1])
                if window_size > max_window_size:
                    max_window_size = window_size

        if max_window_size > 0:
            model_cache_config._max_window_size = max_window_size

        return model_cache_config

    def fetch_cache(
        self,
        request_id: str,
        tokens: list[int],
        extra_keys: tuple[Any, ...] | None = None,
        extra_key_token_start: int | None = None,
        extra_key_ranges: list[tuple[int, tuple[Any, ...]]] | None = None,
    ) -> tuple[BlockTable | None, list[int]]:
        """
        Find cached prefix blocks for the given tokens.

        Args:
            request_id: Unique request identifier
            tokens: Input token sequence
            extra_keys: Additional keys for hash (e.g., VLM image hash)

        Returns:
            Tuple of (block_table, remaining_tokens)
            - block_table: BlockTable if prefix found, None otherwise
            - remaining_tokens: Tokens that need processing
        """
        if not tokens:
            return None, tokens

        # Try to find shared prefix blocks
        shared_block_ids, remaining = self.paged_cache.find_shared_prefix(
            tokens,
            extra_keys=extra_keys,
            extra_key_token_start=extra_key_token_start,
            extra_key_ranges=extra_key_ranges,
        )

        if shared_block_ids:
            # Create block table for this request with shared blocks
            block_table = self.paged_cache.create_block_table(request_id)

            for block_id in shared_block_ids:
                # Increment ref count for sharing
                self.paged_cache.increment_ref(block_id)
                block = self.paged_cache.allocated_blocks.get(block_id)
                if block:
                    block_table.block_ids.append(block_id)
                    block_table.num_tokens += block.token_count

            num_prefix_tokens = len(tokens) - len(remaining)
            self._hits += 1
            self._tokens_saved += num_prefix_tokens
            self._tokens_matched_total += num_prefix_tokens
            self._tokens_requested_total += len(tokens)

            logger.debug(
                f"Cache hit for {request_id}: "
                f"{len(shared_block_ids)} blocks, {num_prefix_tokens} tokens"
            )

            return block_table, remaining

        # Try prefix index for longer matches
        best_match = self._find_best_prefix_match(tokens, extra_keys=extra_keys)
        if best_match:
            prefix_len, matched_block_ids, num_blocks, chain_hashes = best_match

            # Re-validate the stored chain before reuse: index entries can
            # outlive their blocks (eviction, reuse for other content), and
            # handing out a reassigned block would splice foreign KV into
            # this request. Acquire block by block and stop at the first
            # mismatch -- the validated prefix up to that point is still
            # correct. The old code skipped missing blocks but kept the full
            # prefix_len, leaving holes in the block table.
            acquired: list[CacheBlock] = []
            for block_id, expected_hash in zip(
                matched_block_ids[:num_blocks], chain_hashes
            ):
                block = self.paged_cache.acquire_cached_block(block_id, expected_hash)
                if block is None:
                    # Chain broken: self-heal the stale index entry so the
                    # next lookup does not walk the same dead chain.
                    self._prefix_index.pop(chain_hashes[num_blocks - 1], None)
                    break
                acquired.append(block)

            if acquired:
                block_table = self.paged_cache.create_block_table(request_id)
                for block in acquired:
                    block_table.block_ids.append(block.block_id)
                    block_table.num_tokens += block.token_count

                matched_tokens = block_table.num_tokens
                remaining = tokens[matched_tokens:]
                self._hits += 1
                self._tokens_saved += matched_tokens
                self._tokens_matched_total += matched_tokens
                self._tokens_requested_total += len(tokens)

                logger.debug(
                    f"Prefix index hit for {request_id}: "
                    f"{matched_tokens} tokens matched"
                )

                return block_table, remaining

        # No cache hit
        self._misses += 1
        self._tokens_requested_total += len(tokens)
        logger.debug(f"Cache miss for {request_id}")
        return None, tokens

    def store_cache(
        self,
        request_id: str,
        tokens: list[int],
        cache_data: list[Any],
        model_cache_config: ModelCacheConfig | None = None,
        boundary_snapshots: dict[int, list[Any]] | None = None,
        extra_keys: tuple[Any, ...] | None = None,
        extra_key_token_start: int | None = None,
        extra_key_ranges: list[tuple[int, tuple[Any, ...]]] | None = None,
        hot_cache_write_back: bool = True,
    ) -> BlockTable | None:
        """
        Store computed cache for future reuse.

        In paged SSD-only mode, this method:
        1. Allocates block metadata in PagedCacheManager
        2. Extracts tensor slices for each block
        3. Saves each block's data to paged SSD via PagedSSDCacheManager

        Args:
            request_id: Unique request identifier
            tokens: Token sequence that was processed
            cache_data: The computed KV cache to store. Can be:
                - List of KVCache objects (legacy)
                - List of dicts with 'state': (keys, values) tensors (preferred)
            model_cache_config: Optional cache configuration with per-layer type
                information. If None, assumes all layers use KVCache.
            boundary_snapshots: Optional mapping of token_count -> extracted cache
                states for intermediate block boundaries. Used to store per-block
                ArraysCache state instead of placeholders in hybrid models.
            hot_cache_write_back: When False, SSD-backed hot cache is bypassed
                for newly stored dirty blocks.

        Returns:
            BlockTable for the stored cache, or None on failure
        """
        if not tokens:
            return None

        # Check if cache_data contains extracted tensor states
        is_tensor_data = (
            cache_data
            and isinstance(cache_data, list)
            and len(cache_data) > 0
            and isinstance(cache_data[0], dict)
            and "state" in cache_data[0]
        )

        # Extract cache type information for SSD storage
        layer_cache_types = None
        layer_meta_states = None
        if model_cache_config:
            layer_cache_types = model_cache_config.get_type_names()
            # Extract meta_states if available in cache_data
            layer_meta_states = [
                cache_data[i].get("meta_state", ()) if i < len(cache_data) else ()
                for i in range(model_cache_config.num_layers)
            ]
        elif is_tensor_data:
            # Try to extract type info from cache_data itself
            layer_cache_types = [
                # Prefer class_name for TurboQuant (cache_type maps to 'KVCache'),
                # fall back to cache_type for all standard mlx-lm types.
                (
                    layer_state.get(
                        "class_name", layer_state.get("cache_type", "KVCache")
                    )
                    if layer_state.get("class_name", "")
                    in ("TurboQuantKVCache", "BatchTurboQuantKVCache")
                    else layer_state.get("cache_type", "KVCache")
                )
                for layer_state in cache_data
            ]
            layer_meta_states = [
                layer_state.get("meta_state", ()) for layer_state in cache_data
            ]

        # Get or create block table
        block_table = self.paged_cache.get_block_table(request_id)
        if not block_table:
            block_table = self.paged_cache.create_block_table(request_id)

        # Determine tokens we need to cache (not already in block_table)
        existing_tokens = block_table.num_tokens
        new_tokens = tokens[existing_tokens:]

        if not new_tokens:
            # All tokens already cached
            self._last_partial_tokens_skipped = 0
            self._last_tokens_to_next_block = 0
            return block_table

        # Allocate only full blocks (skip partial trailing block).
        # get_computed_blocks() matches full blocks only (floor division),
        # so partial block data is never used during cache lookup.
        # Skipping partial blocks also ensures is_last_block points to
        # the last full block, which is critical for non-sliceable caches
        # (ArraysCache/RotatingKVCache) that use last-block-only storage.
        num_new_blocks = len(new_tokens) // self.block_size
        trailing_partial_tokens = len(new_tokens) % self.block_size
        self._last_partial_tokens_skipped = trailing_partial_tokens
        self._last_tokens_to_next_block = (
            self.block_size - trailing_partial_tokens
            if trailing_partial_tokens > 0
            else 0
        )
        if trailing_partial_tokens > 0:
            self._partial_block_skips += 1
            self._partial_tokens_skipped += trailing_partial_tokens
            logger.debug(
                "Skipping trailing partial block for %s: %s token(s) not persisted "
                "(block_size=%s, needs +%s token(s) to fill next block)",
                request_id,
                trailing_partial_tokens,
                self.block_size,
                self._last_tokens_to_next_block,
            )

        blocks_saved_to_ssd = 0
        # Supersede-on-extend tracking (rotating models only, see below).
        first_new_block_idx: int | None = None
        tip_block_saved = False

        for i in range(num_new_blocks):
            start_idx = i * self.block_size
            end_idx = min(start_idx + self.block_size, len(new_tokens))
            block_tokens = new_tokens[start_idx:end_idx]

            # Token range in the original sequence (accounting for existing tokens)
            global_start = existing_tokens + start_idx
            global_end = existing_tokens + end_idx

            # Compute parent hash for chain-based lookup
            parent_hash = None
            if block_table.block_ids:
                prev_block_id = block_table.block_ids[-1]
                prev_block = self.paged_cache.allocated_blocks.get(prev_block_id)
                if prev_block and prev_block.block_hash:
                    parent_hash = prev_block.block_hash

            block_extra_keys = resolve_block_extra_keys(
                global_end,
                extra_keys=extra_keys,
                extra_key_token_start=extra_key_token_start,
                extra_key_ranges=extra_key_ranges,
            )

            # Check if this block already exists (deduplication)
            if len(block_tokens) == self.block_size:
                existing_block = self.paged_cache.find_cached_block(
                    block_tokens,
                    parent_hash,
                    extra_keys=block_extra_keys,
                )
                if existing_block:
                    # Reuse existing block
                    self.paged_cache.increment_ref(existing_block.block_id)
                    block_table.block_ids.append(existing_block.block_id)
                    block_table.num_tokens += len(block_tokens)
                    continue

            # Allocate new block
            if first_new_block_idx is None:
                first_new_block_idx = len(block_table.block_ids)
            block = self.paged_cache.allocate_block()
            if not block:
                # Handle memory pressure
                if not self.paged_cache.handle_memory_pressure(1):
                    logger.warning(f"Cannot allocate block for {request_id}")
                    break
                block = self.paged_cache.allocate_block()
                if not block:
                    break

            # Set block metadata
            block.token_count = len(block_tokens)
            block_table.block_ids.append(block.block_id)
            block_table.num_tokens += len(block_tokens)

            # Compute chain hash for this block
            block.block_hash = compute_block_hash(
                parent_hash,
                block_tokens,
                extra_keys=block_extra_keys,
                model_name=self.paged_cache.model_name,
            )

            # Register hash for full blocks (for deduplication)
            if len(block_tokens) == self.block_size:
                self.paged_cache.register_block_hash(
                    block, block_tokens, parent_hash, extra_keys=block_extra_keys
                )

            # Extract tensor slice and save to paged SSD
            if is_tensor_data and HAS_MLX and self.paged_ssd_cache is not None:
                cache_seq_len = self._get_cache_seq_len(cache_data)

                # Determine whether extracted cache_data uses:
                # - global indices (full sequence cache, includes reused prefix), or
                # - relative indices (only newly processed suffix).
                #
                # BatchGenerator.extract_cache() currently returns full-sequence cache.
                # When existing_tokens > 0, slicing with relative indices would save
                # wrong KV ranges for new blocks and corrupt future cache hits.
                cache_uses_global_indices = existing_tokens > 0 and cache_seq_len >= (
                    existing_tokens + 1
                )
                if cache_uses_global_indices:
                    cache_start = global_start
                    cache_end = global_end
                else:
                    cache_start = start_idx
                    cache_end = end_idx

                is_last_block = i == num_new_blocks - 1

                # Look up boundary snapshot BEFORE the continuity check.
                # Snapshots are self-contained — they carry the full cache
                # state at this boundary, so the live-cache seq_len gate
                # below does not apply when a snapshot covers this block.
                block_boundary_tc = existing_tokens + end_idx
                snapshot_cache_data = None
                if boundary_snapshots and block_boundary_tc in boundary_snapshots:
                    snapshot_cache_data = boundary_snapshots[block_boundary_tc]

                # Continuity check applies only when we will slice live
                # cache_data for this block. Skipped when:
                #   1. A boundary snapshot exists for this block — snapshots
                #      are self-contained, so the live-cache seq_len gate
                #      does not apply.
                #   2. is_last_block is True — _extract_block_tensor_slice's
                #      last-block branch uses cache_data's full state for
                #      non-sliceable types (RotatingKVCache last window,
                #      CacheList has_valid_state path) and needs no
                #      sliceable seq_len. For sliceable hybrid models the
                #      step-1 path already returns the full prefill length,
                #      so the gate would not fire here anyway.
                if (
                    snapshot_cache_data is None
                    and not is_last_block
                    and cache_seq_len > 0
                    and cache_start >= cache_seq_len
                ):
                    logger.debug(
                        f"Cache continuity broken: cache only has {cache_seq_len} tokens, "
                        f"cannot store block at cache indices [{cache_start}:{cache_end}] "
                        f"(global [{global_start}:{global_end}]). Stopping block allocation."
                    )
                    # Free the block we just allocated (it has no data)
                    self.paged_cache.free_block(block.block_id)
                    block_table.block_ids.pop()
                    block_table.num_tokens -= len(block_tokens)
                    break

                block_kv_data = self._extract_block_tensor_slice(
                    cache_data,
                    cache_start,
                    cache_end,
                    model_cache_config,
                    is_last_block=is_last_block,
                    snapshot_cache_data=snapshot_cache_data,
                )

                if block_kv_data and block.block_hash:
                    # Use per-block meta_states from boundary snapshot when
                    # available. The shared layer_meta_states comes from the
                    # final cache extraction and carries the end-of-request
                    # offset (e.g. 4479) which is wrong for earlier blocks
                    # whose tensor data was captured at an earlier boundary
                    # (e.g. offset=512). Boundary snapshots record the
                    # correct per-boundary meta_state synchronously during
                    # prefill, so we prefer those.
                    block_meta = layer_meta_states
                    if (
                        snapshot_cache_data is not None
                        and layer_meta_states is not None
                    ):
                        per_block = []
                        for lidx in range(len(layer_meta_states)):
                            if (
                                lidx < len(snapshot_cache_data)
                                and isinstance(snapshot_cache_data[lidx], dict)
                                and snapshot_cache_data[lidx].get("meta_state")
                                and snapshot_cache_data[lidx]["meta_state"] != ()
                            ):
                                per_block.append(
                                    snapshot_cache_data[lidx]["meta_state"]
                                )
                            else:
                                per_block.append(layer_meta_states[lidx])
                        block_meta = per_block

                    # Save to paged SSD via PagedSSDCacheManager with cache type info
                    if hot_cache_write_back:
                        saved = self.paged_ssd_cache.save_block(
                            block_hash=block.block_hash,
                            cache_data=block_kv_data,
                            token_count=block.token_count,
                            model_name=self.paged_cache.model_name,
                            layer_cache_types=layer_cache_types,
                            layer_meta_states=block_meta,
                        )
                    else:
                        saved = self.paged_ssd_cache.save_block(
                            block_hash=block.block_hash,
                            cache_data=block_kv_data,
                            token_count=block.token_count,
                            model_name=self.paged_cache.model_name,
                            layer_cache_types=layer_cache_types,
                            layer_meta_states=block_meta,
                            hot_cache_write_back=False,
                        )
                    if saved:
                        blocks_saved_to_ssd += 1
                        if is_last_block:
                            tip_block_saved = True
                        logger.debug(
                            f"Saved block {block.block_id} to tiered cache: "
                            f"tokens [{global_start}:{global_end}], {len(block_kv_data)} layers"
                        )
                    else:
                        logger.warning(
                            f"Failed to save block {block.block_id} to tiered cache"
                        )
                        # Persistence failed: roll back metadata so we don't
                        # retain a block that cannot be reconstructed later.
                        self.paged_cache.free_block(block.block_id)
                        block_table.block_ids.pop()
                        block_table.num_tokens -= len(block_tokens)
                        break
                else:
                    # Failed to extract tensor data - free block and stop
                    logger.debug(
                        f"Failed to extract tensor slice [{global_start}:{global_end}], "
                        f"freeing block {block.block_id} and stopping."
                    )
                    self.paged_cache.free_block(block.block_id)
                    block_table.block_ids.pop()
                    block_table.num_tokens -= len(block_tokens)
                    break

        # Supersede-on-extend: on rotating (sliding-window) models every store
        # of a growing conversation writes one tip block carrying the full
        # sliding-window state of all rotating layers (hundreds of MB fp16 on
        # a gemma3-class model). Restore only ever consumes the newest such
        # block, and the immediate previous tip is kept intact as the
        # walk-back fallback — so the tip two generations back is dead
        # weight. Without stripping it, those blocks fill the hot cache after
        # ~10-20 turns and LRU eviction breaks the prefix chain (multi-turn
        # cache hit collapses to 0%). Steady state after stripping: two heavy
        # blocks per chain.
        if (
            tip_block_saved
            and first_new_block_idx is not None
            and 0 < first_new_block_idx < len(block_table.block_ids)
            and layer_cache_types
            and any(CacheTypeRegistry.is_rotating_family(t) for t in layer_cache_types)
        ):
            prev_tip_id = block_table.block_ids[first_new_block_idx - 1]
            new_tip_id = block_table.block_ids[-1]
            prev_tip = self.paged_cache.allocated_blocks.get(prev_tip_id)
            new_tip = self.paged_cache.allocated_blocks.get(new_tip_id)
            if (
                prev_tip is not None
                and prev_tip.block_hash is not None
                and new_tip is not None
                and new_tip.block_hash is not None
            ):
                superseded = self._rotating_tip_lineage.pop(prev_tip.block_hash, None)
                if superseded is not None:
                    self._strip_rotating_payload(superseded)
                self._rotating_tip_lineage[new_tip.block_hash] = prev_tip.block_hash
                if len(self._rotating_tip_lineage) > _TIP_LINEAGE_MAX_ENTRIES:
                    self._rotating_tip_lineage.clear()

        # Update prefix index
        self._update_prefix_index(tokens, block_table.block_ids, extra_keys=extra_keys)

        # Store entry for request tracking
        self._request_tables[request_id] = BlockCacheEntry(
            block_table=block_table,
            last_access=time.time(),
        )

        logger.debug(
            f"Stored cache for {request_id}: "
            f"{len(block_table.block_ids)} blocks ({blocks_saved_to_ssd} saved to tiered cache), "
            f"{block_table.num_tokens} tokens"
        )

        return block_table

    def _get_cache_seq_len(self, cache_data: list[dict[str, Any]]) -> int:
        """
        Get the sequence length from cache data.

        For hybrid models (e.g., gpt-oss, gemma3 with KVCache + RotatingKVCache layers),
        this finds a standard KVCache layer (full attention) to determine the actual
        seq_len. RotatingKVCache layers use sliding window and have limited seq_len.
        ArraysCache layers don't have a sequence dimension.

        Args:
            cache_data: List of layer states, each containing 'state': (keys, values)

        Returns:
            Sequence length from first sliceable KVCache layer, or max seq_len as fallback
        """
        if not cache_data:
            return 0

        # Non-sliceable cache types use sliding window or have no sequence dimension
        # RotatingKVCache: sliding window, seq_len limited to max_size
        # ArraysCache: no traditional sequence dimension
        non_sliceable_types = {
            "ArraysCache",
            "CacheList",
            "MiniMaxM3BatchKVCache",
        }

        # Step 1: Search for a sliceable KVCache layer (full attention)
        for layer_idx, layer_state in enumerate(cache_data):
            try:
                if "state" not in layer_state:
                    continue

                # Skip non-sliceable cache types (e.g., RotatingKVCache)
                cache_type = layer_state.get("cache_type", "")
                class_name = layer_state.get("class_name", "")
                if (
                    cache_type in non_sliceable_types
                    or class_name in non_sliceable_types
                    or CacheTypeRegistry.is_rotating_family(cache_type)
                    or CacheTypeRegistry.is_rotating_family(class_name)
                ):
                    continue

                state = layer_state["state"]
                keys = state[0] if isinstance(state, (list, tuple)) else state
                # TurboQuant v2: NamedTuple state with .norms attribute
                if hasattr(keys, "norms") and hasattr(keys.norms, "shape"):
                    seq_len = keys.norms.shape[2]
                    logger.debug(
                        f"Found TurboQuantKVCache at layer {layer_idx} with seq_len={seq_len}"
                    )
                    return seq_len
                # TurboQuant v2: SplitState with .low/.high sub-states
                if hasattr(keys, "low") and hasattr(keys.low, "norms"):
                    seq_len = keys.low.norms.shape[2]
                    logger.debug(
                        f"Found TurboQuantKVCache (split) at layer {layer_idx} with seq_len={seq_len}"
                    )
                    return seq_len
                if not hasattr(keys, "shape"):
                    continue

                # KVCache: shape (batch, n_kv_heads, seq_len, head_dim) - 4D
                if len(keys.shape) == 4:
                    seq_len = keys.shape[2]
                    logger.debug(
                        f"Found KVCache at layer {layer_idx} with seq_len={seq_len}"
                    )
                    return seq_len

            except Exception:
                continue

        # Step 2: Fallback - find max seq_len among all 4D tensors
        # This handles pure RotatingKVCache models or unknown cache types.
        # Only skip cache types that do not expose a sequence dimension here.
        # RotatingKVCache must be included because pure RotatingKVCache models
        # have no sliceable KVCache layers for Step 1 to find.
        step2_skip_types = {
            "ArraysCache",
            "CacheList",
            "MiniMaxM3BatchKVCache",
        }
        max_seq_len = 0
        for layer_idx, layer_state in enumerate(cache_data):
            try:
                if "state" not in layer_state:
                    continue
                cache_type = layer_state.get("cache_type", "")
                class_name = layer_state.get("class_name", "")
                if cache_type in step2_skip_types or class_name in step2_skip_types:
                    continue
                state_tuple = layer_state["state"]
                if not isinstance(state_tuple, (list, tuple)) or not state_tuple:
                    continue
                # N-tuple safe: only the first element (the keys-shaped tensor
                # in legacy KVCache, or buf_kv-shaped in PoolingCache) is
                # consulted for seq length here. Caches whose first element is
                # not a 4D KVCache-style tensor naturally skip via the shape
                # check below.
                keys = state_tuple[0]
                if hasattr(keys, "shape") and len(keys.shape) == 4:
                    max_seq_len = max(max_seq_len, keys.shape[2])
            except Exception:
                continue

        if max_seq_len > 0:
            # Normal result for all-non-sliceable-KVCache models
            # (e.g. DeepSeek V4 with RotatingKVCache + PoolingCache).
            # Returns the sliding-window length, not a sequence length —
            # callers that depend on full prefill length (continuity
            # check, reconstruction concat) bypass this layer's
            # contribution via the snapshot / is_last_block paths.
            logger.debug(
                f"Cache seq_len resolved from non-sliceable layer "
                f"(window={max_seq_len})"
            )
            return max_seq_len

        # Step 3: CacheList fallback — check sub-states for seq_len
        # This handles all-CacheList models (e.g., deepseek_v32)
        for layer_state in cache_data:
            if (
                layer_state.get("cache_type") == "CacheList"
                or layer_state.get("class_name") == "CacheList"
            ):
                sub_states = layer_state.get("state", [])
                for sub_state in sub_states:
                    if isinstance(sub_state, (list, tuple)) and len(sub_state) >= 2:
                        sub_keys = sub_state[0]
                        if hasattr(sub_keys, "shape") and len(sub_keys.shape) == 4:
                            seq_len = sub_keys.shape[2]
                            logger.debug(f"Using CacheList sub-cache seq_len={seq_len}")
                            return seq_len

        return 0

    def _strip_rotating_payload(self, block_hash: bytes) -> bool:
        """Replace a superseded tip block's rotating payload with placeholders.

        Sliceable layers (KVCache/TurboQuant slices) in the block are kept —
        only RotatingKVCache-family layer states are replaced with the same
        ``(mx.zeros((1,)), mx.zeros((1,)))`` placeholder that non-tip blocks
        receive at store time, so restore treats the stripped block exactly
        like any other placeholder block (walk-back or reject).

        The rewrite goes through ``forget_block()`` + ``save_block()``: the
        hot-cache entry is removed via ``_hot_cache_remove`` (which also
        forgets the shared-budget accounting) and the slim payload re-enters
        via ``_hot_cache_put`` (which re-registers it), so the hot-cache and
        shared-budget byte counters stay consistent. In SSD mode the slim
        payload is re-enqueued and overwrites the same hash-derived file
        path.

        Returns:
            True if the block was rewritten with at least one layer stripped.
        """
        if self.paged_ssd_cache is None or not HAS_MLX:
            return False
        try:
            data, meta = self.paged_ssd_cache.load_block_with_metadata(block_hash)
            if not data or not meta:
                return False
            types = meta.get("layer_cache_types") or []
            new_data: list[Any] = []
            stripped = 0
            for i, layer in enumerate(data):
                type_name = types[i] if i < len(types) else "KVCache"
                if (
                    CacheTypeRegistry.is_rotating_family(type_name)
                    and isinstance(layer, (list, tuple))
                    and len(layer) >= 2
                    and hasattr(layer[0], "shape")
                    and tuple(layer[0].shape) != (1,)
                ):
                    new_data.append((mx.zeros((1,)), mx.zeros((1,))))
                    stripped += 1
                else:
                    new_data.append(layer)
            if stripped == 0:
                return False
            # save_block dedups on an existing hash, so drop the old entry
            # (hot cache, pending writes, SSD index) first. The brief gap is
            # benign: a concurrent restore either already loaded the old
            # payload or sees the placeholder version, which the
            # walk-back/reject path handles like any partial match.
            self.paged_ssd_cache.forget_block(block_hash)
            saved = self.paged_ssd_cache.save_block(
                block_hash=block_hash,
                cache_data=new_data,
                token_count=int(meta.get("token_count") or self.block_size),
                model_name=meta.get("model_name", self.paged_cache.model_name),
                layer_cache_types=types or None,
                layer_meta_states=meta.get("layer_meta_states"),
            )
            if saved:
                logger.debug(
                    "Stripped rotating payload from superseded tip block %s "
                    "(%d of %d layers)",
                    block_hash.hex()[:16],
                    stripped,
                    len(data),
                )
            return bool(saved)
        except Exception as e:
            logger.debug(
                "Rotating payload strip failed for %s: %s", block_hash.hex()[:16], e
            )
            return False

    def _extract_block_tensor_slice(
        self,
        cache_data: list[dict[str, Any]],
        start_idx: int,
        end_idx: int,
        model_cache_config: ModelCacheConfig | None = None,
        is_last_block: bool = False,
        snapshot_cache_data: list[dict[str, Any]] | None = None,
    ) -> list[tuple[Any, Any]] | None:
        """
        Extract tensor slices for a single block from cache data.

        Supports different cache types (KVCache, RotatingKVCache, ArraysCache)
        with type-aware slicing. For non-sliceable types like ArraysCache,
        returns the full state.

        For RotatingKVCache layers specifically:
        - Last block: stores the full RotatingKVCache state (keys, values)
        - Non-last blocks: stores a placeholder (mx.zeros((1,)), mx.zeros((1,)))
          to preserve layer count while minimizing storage
        - Boundary snapshot: if a snapshot was captured at this block's boundary,
          the snapshot state is used instead of a placeholder

        During restore, a partial prefix match that ends on a placeholder block
        first attempts walk-back truncation to the latest block with valid
        non-sliceable state. If no such block exists, the cache hit is rejected.

        Args:
            cache_data: List of layer states, each containing 'state': (keys, values)
                or other cache-type-specific format
            start_idx: Start token index in the sequence
            end_idx: End token index in the sequence
            model_cache_config: Optional model cache configuration with per-layer
                type information
            is_last_block: If True, this is the last block being stored. For
                RotatingKVCache layers, only the last block stores full state.
            snapshot_cache_data: Optional boundary snapshot cache data for this
                block. When provided, non-sliceable layers use the snapshot state
                instead of a placeholder.

        Returns:
            List of (keys_slice, values_slice) for each layer, or None on failure
        """
        if not HAS_MLX or not cache_data:
            return None

        try:
            block_slices = []
            for layer_idx, layer_state in enumerate(cache_data):
                if "state" not in layer_state:
                    continue

                # Determine cache type for this layer
                cache_type_name = layer_state.get("cache_type", "KVCache")
                if model_cache_config and layer_idx < len(
                    model_cache_config.layer_configs
                ):
                    cache_type_name = model_cache_config.layer_configs[
                        layer_idx
                    ].class_name

                handler = CacheTypeRegistry.get_handler_by_class_name(cache_type_name)

                if cache_type_name in ("TurboQuantKVCache", "BatchTurboQuantKVCache"):
                    # TurboQuant v2: NamedTuple state from mlx-vlm
                    from ..turboquant_kv import _slice_state_range, _state_length

                    state = layer_state["state"]
                    if not isinstance(state, (list, tuple)) or len(state) < 2:
                        block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                        continue
                    k_state, v_state = state[0], state[1]
                    # Unwrap _QuantizedStateProxy if present
                    if hasattr(k_state, "_state"):
                        k_state = k_state._state
                    if hasattr(v_state, "_state"):
                        v_state = v_state._state
                    seq_len = _state_length(k_state)
                    actual_end = min(end_idx, seq_len)
                    if start_idx >= actual_end:
                        block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                        continue
                    ks = _slice_state_range(k_state, start_idx, actual_end)
                    vs = _slice_state_range(v_state, start_idx, actual_end)
                    block_slices.append(
                        (
                            "__turboquant_v2__",
                            (ks, vs),
                        )
                    )
                elif handler.supports_block_slicing:
                    # Standard 4D KV cache slicing
                    state = layer_state["state"]
                    if not isinstance(state, (list, tuple)) or len(state) < 2:
                        # Placeholder from boundary snapshot (skipped sliceable layer).
                        continue

                    axis_info = handler.get_state_axis_info()
                    if len(state) != 2 or len(axis_info) != 2:
                        seq_len = handler.get_state_seq_len_from_tuple(tuple(state))
                        actual_end = min(end_idx, seq_len)
                        if seq_len <= 0 or start_idx >= actual_end:
                            continue

                        sliced_elements = []
                        for info, elem in zip(axis_info, state):
                            if elem is None:
                                sliced_elements.append(None)
                                continue
                            if (
                                info.sliceable
                                and info.sequence_axis is not None
                                and hasattr(elem, "shape")
                                and info.sequence_axis < len(elem.shape)
                            ):
                                slices = [slice(None)] * len(elem.shape)
                                slices[info.sequence_axis] = slice(
                                    start_idx, actual_end
                                )
                                sliced_elements.append(
                                    self._clone_tensor(elem[tuple(slices)])
                                )
                            else:
                                sliced_elements.append(
                                    self._clone_tensor(elem)
                                    if hasattr(elem, "shape")
                                    else elem
                                )
                        block_slices.append(
                            ("__nstate__", cache_type_name, sliced_elements)
                        )
                        continue

                    keys, values = state

                    # KV cache shape: (batch, n_kv_heads, seq_len, head_dim)
                    # Slice along seq_len dimension (axis 2)
                    if not hasattr(keys, "shape") or len(keys.shape) < 4:
                        # Handle 3D case (no batch dimension)
                        if hasattr(keys, "shape") and len(keys.shape) == 3:
                            seq_len = keys.shape[1]  # (n_kv_heads, seq_len, head_dim)
                            actual_end = min(end_idx, seq_len)
                            if start_idx >= actual_end:
                                continue
                            keys_slice = keys[:, start_idx:actual_end, :]
                            values_slice = values[:, start_idx:actual_end, :]
                        else:
                            logger.debug(
                                f"Layer {layer_idx}: unexpected tensor shape for {cache_type_name}"
                            )
                            continue
                    else:
                        seq_len = keys.shape[2]
                        if end_idx > seq_len:
                            logger.debug(
                                f"Block slice [{start_idx}:{end_idx}] exceeds seq_len {seq_len}"
                            )
                            actual_end = min(end_idx, seq_len)
                            if start_idx >= actual_end:
                                continue
                            keys_slice = keys[:, :, start_idx:actual_end, :]
                            values_slice = values[:, :, start_idx:actual_end, :]
                        else:
                            keys_slice = keys[:, :, start_idx:end_idx, :]
                            values_slice = values[:, :, start_idx:end_idx, :]

                    # Detach slices so block-level eviction can free memory
                    block_slices.append(
                        (
                            self._clone_tensor(keys_slice),
                            self._clone_tensor(values_slice),
                        )
                    )
                elif CacheTypeRegistry.is_rotating_family(cache_type_name):
                    # RotatingKVCache: last-block-only or boundary-snapshot strategy
                    has_valid_state = is_last_block or (
                        snapshot_cache_data is not None
                        and layer_idx < len(snapshot_cache_data)
                    )
                    if has_valid_state:
                        # Use snapshot state if available, otherwise use main state
                        if (
                            snapshot_cache_data is not None
                            and layer_idx < len(snapshot_cache_data)
                            and "state" in snapshot_cache_data[layer_idx]
                        ):
                            state = snapshot_cache_data[layer_idx]["state"]
                        else:
                            state = layer_state["state"]
                        if isinstance(state, (list, tuple)) and len(state) >= 2:
                            keys = state[0]
                            values = state[1]
                            block_slices.append(
                                (self._clone_tensor(keys), self._clone_tensor(values))
                            )
                        else:
                            logger.debug(
                                f"Layer {layer_idx}: RotatingKVCache unexpected state format"
                            )
                            block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                    else:
                        # Non-last block without snapshot: store placeholder
                        block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                elif cache_type_name == "CacheList":
                    state = layer_state["state"]  # List[sub_state]
                    sub_class_names = layer_state.get("sub_class_names") or []
                    if not sub_class_names:
                        # Snapshot entries and older extract paths carry the
                        # sub class names only inside the composite
                        # meta_state ([class_names], [sub_meta_states]).
                        meta = layer_state.get("meta_state")
                        if (
                            isinstance(meta, (list, tuple))
                            and len(meta) >= 2
                            and isinstance(meta[0], (list, tuple))
                        ):
                            sub_class_names = [str(n) for n in meta[0]]
                    if not isinstance(state, list) or len(state) == 0:
                        block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                        continue

                    # Check if all sub-caches are sliceable 4D KVCache tensors.
                    # PoolingCache fails this check (its first element buf_kv
                    # is 3D), so a CacheList containing a PoolingCache falls
                    # to the non-sliceable last-block-only branch below.
                    all_sub_sliceable = all(
                        isinstance(ss, (list, tuple))
                        and len(ss) >= 2
                        and hasattr(ss[0], "shape")
                        and len(ss[0].shape) == 4
                        for ss in state
                    )

                    def _sub_class_for(sub_idx):
                        if sub_idx < len(sub_class_names):
                            return sub_class_names[sub_idx]
                        return None

                    def _wrap_sub_marker(sub_idx, elements, storage_class_name=None):
                        # Length-2 element lists round-trip as legacy
                        # ``(keys, values)`` so existing callers (prefix
                        # cache reconstruct, tests) keep their shape. Real
                        # N-tuple sub-states (PoolingCache, BatchKVCache)
                        # surface as ``__nstate__`` markers.
                        if len(elements) == 2:
                            return (elements[0], elements[1])
                        return (
                            "__nstate__",
                            storage_class_name or _sub_class_for(sub_idx),
                            list(elements),
                        )

                    if all_sub_sliceable:
                        # Per-block slicing along sequence axis. Generic over
                        # the full sub-state tuple length so 4D N-tuple caches
                        # (BatchKVCache: 4 elements with offset/padding meta
                        # at indices 2/3) round-trip without dropping
                        # elements past index 1.
                        sub_tensors = []
                        for sub_idx, sub_state in enumerate(state):
                            seq_len = sub_state[0].shape[2]
                            actual_end = min(end_idx, seq_len)
                            sliced_elements = []
                            for elem in sub_state:
                                if (
                                    hasattr(elem, "shape")
                                    and len(elem.shape) == 4
                                    and elem.shape[2] == seq_len
                                ):
                                    if start_idx >= actual_end:
                                        sliced_elements.append(
                                            self._clone_tensor(elem[:, :, 0:0, :])
                                        )
                                    else:
                                        sliced_elements.append(
                                            self._clone_tensor(
                                                elem[:, :, start_idx:actual_end, :]
                                            )
                                        )
                                else:
                                    # Non-sequence element (e.g. BatchKVCache
                                    # offset/left_padding metadata). Pass
                                    # through unsliced.
                                    sliced_elements.append(
                                        self._clone_tensor(elem)
                                        if hasattr(elem, "shape")
                                        else elem
                                    )
                            sub_tensors.append(
                                _wrap_sub_marker(sub_idx, sliced_elements)
                            )
                        block_slices.append(("__cache_list__", sub_tensors))
                    else:
                        # Non-sliceable sub-caches: last-block-only or snapshot.
                        # This is the path PoolingCache takes (3D buf_kv).
                        # Critical fix: clone *all* sub_state elements, not
                        # just the first two, so PoolingCache's third element
                        # `pooled` survives the round-trip. Dropping it was
                        # the V4 cross-session corruption root cause.
                        has_valid_state = is_last_block or (
                            snapshot_cache_data is not None
                            and layer_idx < len(snapshot_cache_data)
                        )
                        if has_valid_state:
                            # Use snapshot if available
                            if (
                                snapshot_cache_data is not None
                                and layer_idx < len(snapshot_cache_data)
                                and "state" in snapshot_cache_data[layer_idx]
                            ):
                                source_layer = snapshot_cache_data[layer_idx]
                                source_state = source_layer["state"]
                            else:
                                source_layer = layer_state
                                source_state = state
                            pooling_delta_ranges = source_layer.get(
                                "pooling_delta_ranges", {}
                            )
                            if isinstance(source_state, list):
                                sub_tensors = []
                                for sub_idx, sub_state in enumerate(source_state):
                                    if (
                                        isinstance(sub_state, (list, tuple))
                                        and len(sub_state) >= 1
                                    ):
                                        cloned = [
                                            (
                                                self._clone_tensor(elem)
                                                if hasattr(elem, "shape")
                                                else elem
                                            )
                                            for elem in sub_state
                                        ]
                                        delta_range = pooling_delta_ranges.get(
                                            str(sub_idx)
                                        )
                                        if (
                                            _sub_class_for(sub_idx) == "PoolingCache"
                                            and isinstance(delta_range, (list, tuple))
                                            and len(delta_range) == 2
                                        ):
                                            cloned.append(
                                                mx.array(delta_range, dtype=mx.int64)
                                            )
                                            sub_tensors.append(
                                                _wrap_sub_marker(
                                                    sub_idx,
                                                    cloned,
                                                    POOLING_CACHE_DELTA_CLASS,
                                                )
                                            )
                                        else:
                                            sub_tensors.append(
                                                _wrap_sub_marker(sub_idx, cloned)
                                            )
                                block_slices.append(("__cache_list__", sub_tensors))
                            else:
                                block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                        else:
                            block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                else:
                    # Other non-sliceable cache (ArraysCache/MambaCache or
                    # model-specific caches such as MiniMax M3). N-tuple
                    # caches keep every element via the SSD V3 marker format.
                    # GDN recurrent state summarizes the ENTIRE sequence in a
                    # fixed-size matrix. Each block boundary snapshot captures
                    # the state at that point in the sequence. Without a snapshot,
                    # non-last blocks get a placeholder so partial matches are
                    # detected and rejected during reconstruction.
                    has_valid_state = is_last_block or (
                        snapshot_cache_data is not None
                        and layer_idx < len(snapshot_cache_data)
                    )
                    if has_valid_state:
                        # Use snapshot state if available, otherwise main state
                        if (
                            snapshot_cache_data is not None
                            and layer_idx < len(snapshot_cache_data)
                            and "state" in snapshot_cache_data[layer_idx]
                        ):
                            state = snapshot_cache_data[layer_idx]["state"]
                        else:
                            state = layer_state["state"]
                        if isinstance(state, (list, tuple)) and len(state) > 2:
                            cloned = [
                                (
                                    self._clone_tensor(elem)
                                    if hasattr(elem, "shape")
                                    else elem
                                )
                                for elem in state
                            ]
                            block_slices.append(("__nstate__", cache_type_name, cloned))
                        elif isinstance(state, (list, tuple)) and len(state) >= 2:
                            conv_state = (
                                state[0] if state[0] is not None else mx.array([])
                            )
                            ssm_state = (
                                state[1] if state[1] is not None else mx.array([])
                            )
                            block_slices.append(
                                (
                                    self._clone_tensor(conv_state),
                                    self._clone_tensor(ssm_state),
                                )
                            )
                        else:
                            logger.debug(
                                f"Layer {layer_idx}: {cache_type_name} unexpected state format"
                            )
                            block_slices.append((mx.zeros((1,)), mx.zeros((1,))))
                    else:
                        # Non-last block without snapshot: store placeholder
                        block_slices.append((mx.zeros((1,)), mx.zeros((1,))))

            return block_slices if block_slices else None

        except Exception as e:
            logger.warning(f"Failed to extract block tensor slice: {e}")
            return None

    @staticmethod
    def _is_placeholder_state(data) -> bool:
        """Check if block layer data is a last-block-only placeholder.

        Non-sliceable cache types (ArraysCache, CacheList with non-sliceable
        sub-caches) store real state only in the last block of each save
        operation. All other blocks get a placeholder: ``(mx.zeros((1,)),
        mx.zeros((1,)))``.

        Returns True if *data* is such a placeholder.
        """
        # CacheList real sub-cache data is stored as a list, never a placeholder
        if isinstance(data, list):
            return False
        if isinstance(data, tuple) and len(data) == 2:
            first = data[0]
            if hasattr(first, "shape") and first.shape == (1,):
                return True
        return False

    def _find_walk_back_truncation_point(
        self,
        all_block_data: list[list[Any]],
        layer_cache_types: list[str] | None,
    ) -> int | None:
        """Find the latest block where all non-sliceable layers have valid state.

        In multi-turn conversations, intermediate blocks can accumulate real
        non-sliceable state (ArraysCache/RotatingKVCache/CacheList) from prior
        save operations while later blocks only have placeholders. This method
        walks backwards from the last loaded block to locate the most recent
        block where **every** non-sliceable layer carries real state.

        Returns:
            0-based block index (inclusive) to truncate to, or ``None`` if
            no truncation is needed (last block already valid) or no valid
            fallback block exists.
        """
        if not all_block_data or not layer_cache_types:
            return None

        num_layers = len(all_block_data[0])
        last_idx = len(all_block_data) - 1

        # Identify "problematic" layers: non-sliceable layer type with
        # placeholder state in the last matched block.
        problematic_layers: list[int] = []
        for layer_idx in range(num_layers):
            cache_type = (
                layer_cache_types[layer_idx]
                if layer_idx < len(layer_cache_types)
                else "KVCache"
            )
            handler = CacheTypeRegistry.get_handler_by_class_name(cache_type)
            if handler.supports_block_slicing:
                continue
            if layer_idx < len(all_block_data[last_idx]):
                if self._is_placeholder_state(all_block_data[last_idx][layer_idx]):
                    problematic_layers.append(layer_idx)

        if not problematic_layers:
            return None  # Last block already has valid state for all layers

        # Walk backwards to find the latest block where ALL problematic
        # layers have real (non-placeholder) state.
        for block_idx in range(last_idx - 1, -1, -1):
            block_data = all_block_data[block_idx]
            all_valid = True
            for layer_idx in problematic_layers:
                if layer_idx < len(block_data) and self._is_placeholder_state(
                    block_data[layer_idx]
                ):
                    all_valid = False
                    break
            if all_valid:
                return block_idx

        return None  # No valid fallback -- fall through to existing rejection

    def _clone_tensor(self, tensor: Any) -> Any:
        """Clone a tensor slice to avoid holding the full backing buffer."""
        try:
            if hasattr(mx, "copy"):
                return mx.copy(tensor)
        except Exception:
            pass

        if hasattr(tensor, "copy"):
            try:
                return tensor.copy()
            except Exception:
                pass

        return mx.array(tensor)

    def _apply_window_padding(
        self,
        matched_blocks: int,
        model_cache_config: ModelCacheConfig | None = None,
    ) -> int:
        """Calculate safe restore limit with window padding for hybrid models.

        For models with RotatingKVCache (sliding window attention), we need to
        ensure the sliding window is fully populated when generation starts.
        This means restoring fewer blocks and reprocessing the padding tokens.

        Example (Gemma3: window_size=1024, block_size=256):
            16 blocks matched -> restore 12 blocks (padding 4 blocks)
            The 4 padding blocks (1024 tokens) will be reprocessed to fill
            the RotatingKVCache sliding window.

        Args:
            matched_blocks: Number of matched cache blocks
            model_cache_config: Model cache configuration

        Returns:
            Number of blocks to actually restore (may be less than matched_blocks)
        """
        if model_cache_config is None or not model_cache_config.has_rotating_layers():
            return matched_blocks

        window_size = model_cache_config.get_max_window_size()
        if window_size <= 0:
            return matched_blocks

        padding_blocks = math.ceil(window_size / self.block_size)
        blocks_to_restore = max(0, matched_blocks - padding_blocks)

        if blocks_to_restore < matched_blocks:
            logger.debug(
                f"Window padding: {matched_blocks} blocks matched, "
                f"restoring {blocks_to_restore} blocks "
                f"(padding {padding_blocks} blocks for window_size={window_size})"
            )

        return blocks_to_restore

    def get_cache_for_generation(
        self,
        request_id: str,
    ) -> tuple[list[Any] | None, bool]:
        """
        Get cache data for generation, loading from paged SSD if needed.

        In paged SSD-only mode, cache data is always loaded from paged SSD via
        reconstruct_cache().

        Args:
            request_id: Request identifier

        Returns:
            Tuple of (cache_data, was_loaded_from_ssd)
        """
        entry = self._request_tables.get(request_id)
        if not entry:
            return None, False

        # Get blocks with COW
        _, was_copied = self.paged_cache.get_blocks_for_generation(entry.block_table)

        # In paged SSD-only mode, always reconstruct from paged SSD
        cache_data = self.reconstruct_cache(entry.block_table)
        if cache_data is None:
            return None, False

        entry.last_access = time.time()
        return cache_data, True

    def release_cache(self, request_id: str) -> None:
        """
        Release cache blocks for a completed request.

        Args:
            request_id: Request identifier
        """
        entry = self._request_tables.pop(request_id, None)
        if entry:
            self.paged_cache.delete_block_table(request_id)
            logger.debug(f"Released cache for {request_id}")

    def clear_request_entry(self, request_id: str) -> None:
        """
        Clear request entry from tracking without freeing blocks.

        This removes the request from _request_tables but keeps the cached
        blocks available for prefix matching. Use this after store_cache()
        when the request is complete but cache should remain for future reuse.

        Args:
            request_id: Request identifier
        """
        entry = self._request_tables.pop(request_id, None)
        if entry:
            logger.debug(f"Cleared request entry for {request_id} (blocks retained)")

    def fork_cache(
        self,
        source_request_id: str,
        new_request_id: str,
    ) -> BlockTable | None:
        """
        Fork cache from one request to another (COW).

        In paged SSD-only mode, cache data is always on paged SSD, so we just
        increment reference counts for the blocks.

        Args:
            source_request_id: Source request ID
            new_request_id: New request ID

        Returns:
            Forked BlockTable, or None if source not found
        """
        source_entry = self._request_tables.get(source_request_id)
        if not source_entry:
            return None

        # Fork block table (increments ref counts)
        forked_table = self.paged_cache.fork_block_table(
            source_entry.block_table,
            new_request_id,
        )

        # Create new entry (cache data is on paged SSD)
        self._request_tables[new_request_id] = BlockCacheEntry(
            block_table=forked_table,
            last_access=time.time(),
        )

        logger.debug(f"Forked cache: {source_request_id} -> {new_request_id}")

        return forked_table

    def preload_blocks(self, block_table: BlockTable) -> int:
        """
        Pre-load matched blocks from SSD into hot cache in parallel.

        Call this between fetch_cache() and reconstruct_cache() to
        convert cold-SSD reads into hot-cache hits. Warm-start requests
        (blocks already in hot cache) return 0 with no I/O.

        Args:
            block_table: BlockTable from fetch_cache() containing matched block IDs.

        Returns:
            Number of blocks successfully preloaded into hot cache.
        """
        if self.paged_ssd_cache is None:
            return 0
        if not block_table or not block_table.block_ids:
            return 0

        block_hashes = []
        for block_id in block_table.block_ids:
            block = self.paged_cache.allocated_blocks.get(block_id)
            if block and block.block_hash is not None:
                block_hashes.append(block.block_hash)

        if not block_hashes:
            return 0

        return self.paged_ssd_cache.preload_matched_blocks(block_hashes)

    def reconstruct_cache(
        self,
        block_table: BlockTable,
        promote_to_hot_cache: bool = True,
    ) -> list[Any] | None:
        """
        Reconstruct cache objects from paged SSD-stored block data.

        This method supports multiple cache types (KVCache, RotatingKVCache,
        ArraysCache) and uses stored type information for proper reconstruction.

        In paged SSD-only mode, this method:
        1. Loads block tensor data from paged SSD via PagedSSDCacheManager
        2. Gets cache type info from paged SSD metadata
        3. Concatenates tensors for each layer (or uses full state for non-sliceable)
        4. Creates appropriate cache objects for inference

        If some blocks cannot be loaded, this method will use only the valid
        prefix blocks and update block_table in-place.

        Args:
            block_table: BlockTable containing block IDs to reconstruct from.
                Will be modified in-place if partial reconstruction.
            promote_to_hot_cache: When False, SSD-loaded blocks are not retained
                in hot cache after active KV reconstruction.

        Returns:
            List of reconstructed cache objects (one per layer),
            or None if reconstruction fails completely
        """
        if not block_table or not block_table.block_ids:
            return None

        if not HAS_MLX:
            logger.warning("Cannot reconstruct cache: MLX not available")
            return None

        if self.paged_ssd_cache is None:
            logger.warning(
                "Cannot reconstruct cache: PagedSSDCacheManager not configured"
            )
            return None

        try:
            # Collect cache data from valid blocks (stop at first invalid)
            all_block_data = []
            valid_block_count = 0
            valid_token_count = 0

            # Cache type information from blocks.
            # Anchor the per-block comparison to the live model's signature
            # rather than block 0's metadata. Bootstrapping from block 0
            # means a stale block (saved before TurboQuant/MTP toggled)
            # silently becomes the truth and every newer, correctly-typed
            # block trips the mismatch warning forever. With the live
            # signature as the reference, the stale block 0 itself gets
            # forgotten on the first failed comparison and reuse can
            # extend past it on the next request.
            manager_expected = (
                getattr(self.paged_ssd_cache, "_expected_layer_cache_types", None)
                if self.paged_ssd_cache is not None
                else None
            )
            # Only adopt when the manager really has a concrete signature.
            # Guard against MagicMock auto-attrs in tests and empty lists
            # (which would unanimously trip the mismatch check); fall back
            # to the historical block-0 bootstrap in those cases.
            if isinstance(manager_expected, (list, tuple)) and manager_expected:
                layer_cache_types = list(manager_expected)
            else:
                layer_cache_types = None
            first_block_meta_states = None  # meta_states from first block
            last_block_meta_states = (
                None  # meta_states from last block (for non-sliceable caches)
            )
            all_block_meta_states = []  # per-block meta_states for walk-back truncation

            for idx, block_id in enumerate(block_table.block_ids):
                block = self.paged_cache.allocated_blocks.get(block_id)
                if not block:
                    logger.debug(
                        f"Block {block_id} not found, using {valid_block_count} "
                        f"valid blocks ({valid_token_count} tokens)"
                    )
                    break  # Stop at first missing block, use valid prefix

                # Load block data from paged SSD
                if block.block_hash is None:
                    logger.debug(
                        f"Block {block_id} has no block_hash, "
                        f"using {valid_block_count} valid blocks"
                    )
                    break  # Stop here, use valid prefix

                # Load with metadata for type information
                if promote_to_hot_cache:
                    block_data, block_metadata = (
                        self.paged_ssd_cache.load_block_with_metadata(block.block_hash)
                    )
                else:
                    block_data, block_metadata = (
                        self.paged_ssd_cache.load_block_with_metadata(
                            block.block_hash,
                            promote_to_hot_cache=False,
                        )
                    )
                if block_data is None:
                    logger.debug(
                        f"Failed to load block {block_id} from tiered cache, "
                        f"using {valid_block_count} valid blocks"
                    )
                    # Remove failed block from hash cache to prevent future false hits
                    if block.block_hash is not None:
                        self.paged_cache.cached_block_hash_to_block.pop(
                            block.block_hash, block.block_id
                        )
                        self.paged_cache._notify_hash_dropped(block.block_hash)
                        logger.debug(
                            f"Removed missing block {block_id} from hash cache"
                        )
                    break  # Stop here, use valid prefix

                # Validate model_name to prevent cross-model cache contamination
                if block_metadata:
                    block_model_name = block_metadata.get("model_name", "")
                    current_model_name = self.paged_cache.model_name

                    # If current model has a name, validate against block's model
                    if current_model_name:
                        if not block_model_name:
                            # Block was saved without model_name (old cache), skip it
                            logger.warning(
                                f"Block has no model_name (legacy cache), "
                                f"current model is '{current_model_name}'. Invalidating cache hit."
                            )
                            self._forget_incompatible_ssd_block(block.block_hash)
                            break  # Stop here, don't use this block
                        elif block_model_name != current_model_name:
                            # Block is from a different model
                            logger.warning(
                                f"Cache model mismatch: block is for '{block_model_name}', "
                                f"current model is '{current_model_name}'. Invalidating cache hit."
                            )
                            self._forget_incompatible_ssd_block(block.block_hash)
                            break  # Stop here, don't use this block

                    # Validate num_layers to catch cross-model cache issues
                    block_num_layers = block_metadata.get("num_layers", 0)
                    if self.expected_num_layers > 0 and block_num_layers > 0:
                        if block_num_layers != self.expected_num_layers:
                            logger.warning(
                                f"Cache layer count mismatch: block has {block_num_layers} layers, "
                                f"expected {self.expected_num_layers}. Invalidating cache hit."
                            )
                            self._forget_incompatible_ssd_block(block.block_hash)
                            break  # Stop here, don't use this block

                    if "block_size" in block_metadata:
                        block_size = block_metadata.get("block_size", 0)
                        if block_size and block_size != self.block_size:
                            logger.warning(
                                f"Cache block size mismatch: block has block_size={block_size}, "
                                f"expected {self.block_size}. Invalidating cache hit."
                            )
                            self._forget_incompatible_ssd_block(block.block_hash)
                            break  # Stop here, don't use this block
                        if not block_size and self.block_size:
                            logger.warning(
                                "Block has no block_size metadata (legacy cache), "
                                f"current block_size is {self.block_size}. Invalidating cache hit."
                            )
                            self._forget_incompatible_ssd_block(block.block_hash)
                            break  # Stop here, don't use this block

                # Extract type info from block metadata
                if block_metadata:
                    block_layer_cache_types = block_metadata.get("layer_cache_types")
                    if layer_cache_types is None:
                        layer_cache_types = block_layer_cache_types
                    elif (
                        block_layer_cache_types is not None
                        and self._canonical_layer_cache_types(block_layer_cache_types)
                        != self._canonical_layer_cache_types(layer_cache_types)
                    ):
                        logger.warning(
                            "Cache layer type mismatch at block %s: got %s, "
                            "expected %s. Truncating cached prefix before this "
                            "block.",
                            block_id,
                            block_layer_cache_types,
                            layer_cache_types,
                        )
                        self._forget_incompatible_ssd_block(
                            block.block_hash, block.block_id
                        )
                        break

                    # Expectation-gated signature fields (TurboQuant depth,
                    # CacheList sub composition). Hot-cache / pending-write
                    # loads bypass the manager's index-scan compatibility
                    # check, so the restore loop must gate them here.
                    signature_gate = getattr(
                        self.paged_ssd_cache, "is_signature_compatible", None
                    )
                    if callable(signature_gate) and not signature_gate(
                        block_metadata.get("cache_signature", "")
                    ):
                        logger.warning(
                            "Cache signature mismatch at block %s "
                            "(TurboQuant depth or CacheList sub composition). "
                            "Truncating cached prefix before this block.",
                            block_id,
                        )
                        self._forget_incompatible_ssd_block(
                            block.block_hash, block.block_id
                        )
                        break

                    # Track meta_states from first and last blocks
                    # Non-sliceable caches (RotatingKVCache) need last block's meta_state
                    block_layer_meta_states = block_metadata.get("layer_meta_states")
                    if first_block_meta_states is None:
                        first_block_meta_states = block_layer_meta_states
                    # Always update last to track the most recent
                    last_block_meta_states = block_layer_meta_states
                    all_block_meta_states.append(block_layer_meta_states)
                else:
                    # A block that loads with data but no metadata at all
                    # cannot be trusted: it skipped every per-block validation
                    # gate above (model_name, num_layers, block_size, layer
                    # types), and continuing would leave the first/last
                    # meta_state trackers stale (pairing this block's tensors
                    # with an earlier block's meta). Truncate the chain here
                    # and use the valid prefix, like the load-failure path.
                    logger.warning(
                        f"Block {block_id} loaded without metadata. "
                        f"Truncating cached prefix before this block "
                        f"({valid_block_count} valid blocks)."
                    )
                    self._forget_incompatible_ssd_block(
                        block.block_hash, block.block_id
                    )
                    break

                # Validate loaded data (pass cache types for hybrid models)
                if not self._validate_block_cache_data(block_data, layer_cache_types):
                    logger.debug(
                        f"Block {block_id} has invalid layer data from tiered cache, "
                        f"using {valid_block_count} valid blocks"
                    )
                    break  # Stop here, use valid prefix

                all_block_data.append(block_data)
                valid_block_count += 1
                valid_token_count += block.token_count

            # If we have fewer valid blocks than requested, update block_table
            if valid_block_count < len(block_table.block_ids):
                if valid_block_count == 0:
                    # Free ref_counts for all blocks before returning
                    for bid in block_table.block_ids:
                        self.paged_cache.free_block(bid)
                    block_table.block_ids.clear()
                    block_table.num_tokens = 0
                    return None  # No valid blocks at all

                # Free ref_counts for blocks we are about to drop
                for bid in block_table.block_ids[valid_block_count:]:
                    self.paged_cache.free_block(bid)

                # Truncate block_table to valid prefix
                original_blocks = len(block_table.block_ids)
                block_table.block_ids = block_table.block_ids[:valid_block_count]
                block_table.num_tokens = valid_token_count
                logger.info(
                    f"Partial cache reconstruction: {valid_block_count}/{original_blocks} "
                    f"blocks, {valid_token_count} tokens"
                )

            if not all_block_data:
                return None

            # Get number of layers from first block
            num_layers = len(all_block_data[0])
            if num_layers == 0:
                return None

            # --- Pre-scan: walk-back truncation for non-sliceable caches ---
            # If the last loaded block has a placeholder for any non-sliceable
            # layer (ArraysCache/RotatingKVCache/non-sliceable CacheList), walk
            # backwards to find the latest block where ALL such layers carry
            # real state. This recovers intermediate blocks from multi-turn
            # conversations instead of rejecting the entire cache.
            if all_block_data and layer_cache_types:
                trunc_idx = self._find_walk_back_truncation_point(
                    all_block_data, layer_cache_types
                )
                if trunc_idx is not None:
                    new_count = trunc_idx + 1
                    dropped_count = len(all_block_data) - new_count

                    # Free ref_counts for dropped blocks
                    for bid in block_table.block_ids[new_count:]:
                        self.paged_cache.free_block(bid)

                    # Truncate data structures
                    all_block_data = all_block_data[:new_count]
                    block_table.block_ids = block_table.block_ids[:new_count]
                    valid_token_count = sum(
                        self.paged_cache.allocated_blocks[bid].token_count
                        for bid in block_table.block_ids
                        if bid in self.paged_cache.allocated_blocks
                    )
                    block_table.num_tokens = valid_token_count

                    # Update meta_states to the truncation-point block
                    if trunc_idx < len(all_block_meta_states):
                        last_block_meta_states = all_block_meta_states[trunc_idx]

                    logger.info(
                        f"Walk-back truncation: dropped {dropped_count} trailing "
                        f"block(s) with placeholder non-sliceable state, keeping "
                        f"{new_count} block(s) ({valid_token_count} tokens)"
                    )

            # Reconstruct caches for each layer
            reconstructed_caches = []
            healed_tq_layers = 0

            for layer_idx in range(num_layers):
                # Determine cache type for this layer
                cache_type_name = "KVCache"
                if layer_cache_types and layer_idx < len(layer_cache_types):
                    cache_type_name = layer_cache_types[layer_idx]

                handler = CacheTypeRegistry.get_handler_by_class_name(cache_type_name)

                # === CacheList: dedicated branch (before standard 2-tuple unpack) ===
                if cache_type_name == "CacheList":
                    last_block_layer_data = all_block_data[-1][layer_idx]

                    # Placeholder detection (partial match → reject for
                    # non-sliceable CacheList, e.g. containing ArraysCache)
                    if (
                        isinstance(last_block_layer_data, tuple)
                        and len(last_block_layer_data) == 2
                        and hasattr(last_block_layer_data[0], "shape")
                        and last_block_layer_data[0].shape == (1,)
                    ):
                        logger.info(
                            f"CacheList layer {layer_idx}: partial prefix match "
                            f"detected (placeholder). Rejecting cache."
                        )
                        return None

                    # Each sub_state in block_data may be either:
                    # - a legacy 2-tuple ``(keys, values)``, or
                    # - an ``('__nstate__', class_name, [elements])`` marker
                    #   emitted by the N-tuple-aware extract path (preserves
                    #   PoolingCache's full 3-tuple state).
                    # _sub_state_elements normalizes both to a raw element
                    # list so downstream concat / unpack does not have to
                    # branch on marker shape.
                    def _sub_state_elements(sub_state):
                        if (
                            isinstance(sub_state, tuple)
                            and len(sub_state) >= 3
                            and isinstance(sub_state[0], str)
                            and sub_state[0] == "__nstate__"
                        ):
                            return list(sub_state[2])
                        if isinstance(sub_state, (list, tuple)) and len(sub_state) >= 1:
                            return list(sub_state)
                        return None

                    def _sub_state_class(sub_state):
                        if (
                            isinstance(sub_state, tuple)
                            and len(sub_state) >= 3
                            and sub_state[0] == "__nstate__"
                        ):
                            return sub_state[1]
                        return None

                    # Collect CacheList data from all blocks that have List[sub_state]
                    cl_block_data = []
                    for block_data in all_block_data:
                        bd = block_data[layer_idx]
                        if isinstance(bd, list) and all(
                            _sub_state_elements(t) is not None for t in bd
                        ):
                            cl_block_data.append(bd)

                    if not cl_block_data:
                        logger.error(
                            f"CacheList layer {layer_idx}: no valid block data found"
                        )
                        return None

                    # Determine sub-cache count from first valid block
                    num_sub_caches = len(cl_block_data[0])

                    # Per-sub-cache class dispatch: sliceable sub-caches
                    # (KVCache) concatenate per-block slices into the full
                    # sequence; non-sliceable sub-caches (RotatingKVCache,
                    # PoolingCache, ArraysCache, BatchPoolingCache) keep
                    # the last block's full state, since each saved block
                    # already snapshots the cache up to its boundary.
                    sub_class_names_for_layer: list[str] = []
                    if (
                        last_block_meta_states
                        and layer_idx < len(last_block_meta_states)
                        and isinstance(last_block_meta_states[layer_idx], (list, tuple))
                        and len(last_block_meta_states[layer_idx]) >= 1
                        and isinstance(
                            last_block_meta_states[layer_idx][0], (list, tuple)
                        )
                    ):
                        sub_class_names_for_layer = list(
                            last_block_meta_states[layer_idx][0]
                        )

                    non_sliceable_sub_classes = {
                        "PoolingCache",
                        "ArraysCache",
                        "SizedArraysCache",
                        "BatchPoolingCache",
                    }

                    def _is_non_sliceable_sub_class(class_name: str) -> bool:
                        return (
                            class_name in non_sliceable_sub_classes
                            or CacheTypeRegistry.is_rotating_family(class_name)
                        )

                    # The store path (_extract_block_tensor_slice) picks ONE
                    # storage mode for the whole layer: per-block slices only
                    # when every sub-state is a >=2-element tuple of 4D
                    # sequence tensors and no sub is a known non-sliceable
                    # class; otherwise EVERY block stores the full cumulative
                    # state of ALL subs at that block's boundary. Restore
                    # mirrors that layer-level decision here. Dispatching per
                    # sub (concat KVCache subs, last-block the rest) silently
                    # concatenated a mixed CacheList's cumulative KV snapshots
                    # into a duplicated sequence (e.g. inkling-style
                    # CacheList(KVCache, ArraysCache): 4+8+12 tokens instead
                    # of 12).
                    any_non_sliceable_sub = any(
                        _is_non_sliceable_sub_class(
                            sub_class_names_for_layer[j]
                            if j < len(sub_class_names_for_layer)
                            else ""
                        )
                        for j in range(num_sub_caches)
                    )
                    last_block_elements = [
                        _sub_state_elements(s) for s in cl_block_data[-1]
                    ]
                    stored_slice_mode = not any_non_sliceable_sub and all(
                        elems is not None
                        and len(elems) >= 2
                        and hasattr(elems[0], "shape")
                        and len(elems[0].shape) == 4
                        for elems in last_block_elements
                    )

                    if len(cl_block_data) > 1 and stored_slice_mode:
                        # Per-block slices: concatenate every sub along the
                        # sequence axis into the full sequence.
                        concatenated_sub_states = []
                        for j in range(num_sub_caches):
                            per_block_elements = [
                                _sub_state_elements(bd[j]) for bd in cl_block_data
                            ]
                            num_elems = len(per_block_elements[-1])
                            cat_elements = []
                            for k in range(num_elems):
                                # Concat sequence-axis tensors; non-sequence
                                # elements (axis-2 mismatch or scalars) take
                                # the last block's value.
                                column = [pb[k] for pb in per_block_elements]
                                first = column[0]
                                if (
                                    hasattr(first, "shape")
                                    and len(first.shape) >= 3
                                    and all(
                                        hasattr(c, "shape")
                                        and c.shape[:2] == first.shape[:2]
                                        for c in column
                                    )
                                ):
                                    if any(d == 0 for d in first.shape):
                                        shape = list(first.shape)
                                        shape[2] = sum(c.shape[2] for c in column)
                                        cat_elements.append(mx.zeros(tuple(shape)))
                                    else:
                                        cat_elements.append(
                                            mx.concatenate(column, axis=2)
                                        )
                                else:
                                    cat_elements.append(column[-1])
                            concatenated_sub_states.append(tuple(cat_elements))
                    else:
                        # Cumulative snapshots (or a single block): the last
                        # matched block already holds the complete state of
                        # every sub at its boundary. PoolingCache V4 markers
                        # are the exception: their append-only pooled tensor
                        # is stored as an absolute-range delta per block and
                        # must be rebuilt in order. A legacy full snapshot may
                        # appear before V4 deltas in an existing cache chain;
                        # it becomes the reconstruction base.
                        concatenated_sub_states = []
                        for j, last_elements in enumerate(last_block_elements):
                            has_pooling_deltas = any(
                                _sub_state_class(bd[j]) == POOLING_CACHE_DELTA_CLASS
                                for bd in cl_block_data
                            )
                            if not has_pooling_deltas:
                                concatenated_sub_states.append(tuple(last_elements))
                                continue

                            pooled_parts = []
                            pooled_length = 0
                            last_pooling_state = None
                            valid_pooling_chain = True
                            for block_cache_list in cl_block_data:
                                sub_state = block_cache_list[j]
                                elements = _sub_state_elements(sub_state)
                                marker_class = _sub_state_class(sub_state)
                                if elements is None:
                                    valid_pooling_chain = False
                                    break

                                if marker_class != POOLING_CACHE_DELTA_CLASS:
                                    if len(elements) < 3:
                                        valid_pooling_chain = False
                                        break
                                    if elements[2] is None:
                                        pooled_parts = []
                                        pooled_length = 0
                                    elif (
                                        hasattr(elements[2], "shape")
                                        and len(elements[2].shape) >= 2
                                    ):
                                        pooled_parts = [elements[2]]
                                        pooled_length = int(elements[2].shape[1])
                                    else:
                                        valid_pooling_chain = False
                                        break
                                    last_pooling_state = list(elements)
                                    continue

                                if (
                                    len(elements) not in (4, 6)
                                    or not hasattr(elements[2], "shape")
                                    or len(elements[2].shape) < 2
                                    or not hasattr(elements[-1], "tolist")
                                ):
                                    valid_pooling_chain = False
                                    break
                                try:
                                    delta_range = elements[-1].tolist()
                                    delta_start, delta_end = (
                                        int(delta_range[0]),
                                        int(delta_range[1]),
                                    )
                                except (IndexError, TypeError, ValueError):
                                    valid_pooling_chain = False
                                    break
                                if (
                                    delta_start != pooled_length
                                    or delta_end < delta_start
                                    or int(elements[2].shape[1])
                                    != delta_end - delta_start
                                ):
                                    valid_pooling_chain = False
                                    break
                                pooled_parts.append(elements[2])
                                pooled_length = delta_end
                                last_pooling_state = list(elements[:-1])

                            if (
                                not valid_pooling_chain
                                or not pooled_parts
                                or last_pooling_state is None
                            ):
                                logger.info(
                                    "CacheList layer %d sub-cache %d: invalid "
                                    "PoolingCache delta chain. Rejecting cache.",
                                    layer_idx,
                                    j,
                                )
                                return None

                            pooled = (
                                pooled_parts[0]
                                if len(pooled_parts) == 1
                                else mx.concatenate(pooled_parts, axis=1)
                            )
                            last_pooling_state[2] = pooled
                            concatenated_sub_states.append(tuple(last_pooling_state))

                    # Build meta_state with correct offsets for reconstructed
                    # sequence length (may differ from original if partial match)
                    meta_state = None
                    if last_block_meta_states and layer_idx < len(
                        last_block_meta_states
                    ):
                        meta_state = last_block_meta_states[layer_idx]

                    if (
                        meta_state
                        and isinstance(meta_state, (list, tuple))
                        and len(meta_state) >= 2
                    ):
                        # Adjust sub-cache offsets to actual concatenated seq_len.
                        # Sliceable sub-caches (KVCache) need offset replaced
                        # with the post-concat seq_len. Non-sliceable
                        # sub-caches (RotatingKVCache, PoolingCache, ...)
                        # keep their last-block meta intact — the sliding
                        # window offset / pool length already encode the
                        # boundary state from the original snapshot.
                        class_names = meta_state[0]
                        adjusted_sub_metas = []
                        for j in range(num_sub_caches):
                            orig_sub_meta = (
                                meta_state[1][j] if j < len(meta_state[1]) else ""
                            )
                            sub_class = (
                                sub_class_names_for_layer[j]
                                if j < len(sub_class_names_for_layer)
                                else ""
                            )
                            if _is_non_sliceable_sub_class(sub_class):
                                adjusted_sub_metas.append(
                                    orig_sub_meta if orig_sub_meta else ""
                                )
                                continue
                            sub_elements = concatenated_sub_states[j]
                            actual_seq_len = None
                            if (
                                sub_elements
                                and len(sub_elements) > 0
                                and hasattr(sub_elements[0], "shape")
                                and len(sub_elements[0].shape) >= 3
                            ):
                                actual_seq_len = sub_elements[0].shape[2]
                            if (
                                actual_seq_len is not None
                                and isinstance(orig_sub_meta, (list, tuple))
                                and len(orig_sub_meta) > 0
                            ):
                                adjusted_sub_metas.append(
                                    (actual_seq_len,) + tuple(orig_sub_meta[1:])
                                )
                            else:
                                adjusted_sub_metas.append(
                                    orig_sub_meta if orig_sub_meta else ""
                                )
                        meta_state = (class_names, adjusted_sub_metas)

                    cache = handler.reconstruct_cache(
                        {"sub_states": concatenated_sub_states}, meta_state
                    )
                    if cache is None:
                        logger.error(
                            f"CacheList layer {layer_idx}: reconstruction failed"
                        )
                        return None
                    reconstructed_caches.append(cache)
                    continue

                # === TurboQuant KV: payload-driven per-block handling ===
                # One block chain can mix payload formats: blocks stored
                # while TurboQuant KV conversion was active carry tagged
                # ('__turboquant_v2__', (ks, vs)) NamedTuple states, while
                # blocks stored without the conversion (chains written by
                # versions that skipped it on chunked-prefill completion,
                # or stored under a different TQ setting) carry plain dense
                # (keys, values) tensors in the SAME chain via dedup.
                # The per-block layer_cache_types mismatch check above
                # truncates such chains when every block's type metadata is
                # present and accurate, but block metadata is optional at
                # load time (missing/corrupt metadata files, failed
                # layer_cache_types round-trips, chains typed from a later
                # block when the first block lacks types) — so the payload
                # itself is the ground truth. Each block must be dispatched
                # on its own payload format: typing the whole chain from
                # the chain-level metadata feeds dense arrays into
                # TurboQuant _concat_state (AttributeError: 'array' object
                # has no attribute 'norms') and rejects the hit. The
                # reverse mixing direction (dense-typed chain with
                # TQ-tagged blocks appended later) is routed here by payload
                # scan for the same reason.
                if cache_type_name in (
                    "TurboQuantKVCache",
                    "BatchTurboQuantKVCache",
                ) or (
                    handler.supports_block_slicing
                    and self._layer_has_turboquant_payload(all_block_data, layer_idx)
                ):
                    entries: list[tuple[int, str, Any, Any]] = []
                    for block_idx, block_data in enumerate(all_block_data):
                        if layer_idx >= len(block_data):
                            continue
                        bd = block_data[layer_idx]
                        if not (isinstance(bd, tuple) and len(bd) == 2):
                            logger.warning(
                                f"TQ layer {layer_idx}: block {block_idx} has "
                                f"unsupported payload type "
                                f"{type(bd).__name__}. Rejecting cache hit."
                            )
                            return None
                        if isinstance(bd[0], str) and bd[0] == "__turboquant_v2__":
                            ks, vs = bd[1]
                            entries.append((block_idx, "tq", ks, vs))
                        elif self._is_placeholder_state(bd):
                            # Empty-slice placeholder written by the TQ store
                            # path: this block carries no KV for the layer,
                            # so the chain cannot be reconstructed.
                            logger.info(
                                f"TQ layer {layer_idx}: block {block_idx} is "
                                f"an empty-slice placeholder. Rejecting cache "
                                f"hit; request will reprocess from scratch."
                            )
                            return None
                        elif (
                            hasattr(bd[0], "shape")
                            and len(bd[0].shape) == 4
                            and hasattr(bd[1], "shape")
                        ):
                            # Plain dense KV slice stored without TQ
                            # conversion (e.g. chunked-prefill completion).
                            entries.append((block_idx, "plain", bd[0], bd[1]))
                        else:
                            logger.warning(
                                f"TQ layer {layer_idx}: block {block_idx} "
                                f"payload is neither a tagged TurboQuant "
                                f"state nor a dense KV slice. Rejecting "
                                f"cache hit."
                            )
                            return None
                    if not entries:
                        logger.debug(f"TQ layer {layer_idx}: no block data")
                        return None

                    # Group consecutive TQ blocks that share (bits, seed) so
                    # each run concatenates in quantized form and dequantizes
                    # once. A homogeneous TQ chain stays a single run — one
                    # concat + one codec rebuild, restored as a quantized
                    # TurboQuantKVCache (upstream #1842: restored long-context
                    # prefixes stay quantized, avoiding full-state
                    # materialization). Mixed chains (plain dense blocks, or
                    # runs with differing (bits, seed)) are dequantized per run
                    # and merged into a dense KVCache, cast back to the dense
                    # blocks' stored dtype: dequantize() emits float32, and an
                    # uncast fp32 layer would silently promote the whole merged
                    # batch cache (2x KV memory) on servers where the TurboQuant
                    # requantize epilogue does not run (TQ disabled, skip_last).
                    # The healed request trades TurboQuant's memory savings for
                    # that one restored prefix.
                    try:
                        from mlx_lm.models.cache import KVCache

                        groups: list[tuple[str, Any, Any]] = []
                        for block_idx, kind, part_k, part_v in entries:
                            if kind == "plain":
                                groups.append(("plain", part_k, part_v))
                                continue
                            params = self._tq_block_params(
                                all_block_meta_states,
                                first_block_meta_states,
                                block_idx,
                                layer_idx,
                            )
                            if params is None:
                                logger.warning(
                                    f"TQ layer {layer_idx}: block {block_idx} "
                                    f"has no TurboQuant (bits, seed) metadata "
                                    f"and no server-configured KV bit depth "
                                    f"is available. Rejecting cache hit "
                                    f"instead of guessing the codec width."
                                )
                                return None
                            if (
                                groups
                                and groups[-1][0] == "tq"
                                and groups[-1][2] == params
                            ):
                                groups[-1][1].append((part_k, part_v))
                            else:
                                groups.append(("tq", [(part_k, part_v)], params))

                        # Homogeneous TQ chain (single quantized run, no dense
                        # blocks): keep it quantized end-to-end like the
                        # pre-payload-driven path, so restored prefixes don't
                        # materialize full state.
                        if len(groups) == 1 and groups[0][0] == "tq":
                            from mlx_vlm.turboquant import TurboQuantKVCache

                            from ..turboquant_kv import (
                                _concat_state_token_axis,
                                _rebuild_codecs,
                                _state_length,
                            )

                            tq_bits, tq_seed = groups[0][2]
                            run_states = groups[0][1]
                            cat_ks = _concat_state_token_axis(
                                [ks for ks, _ in run_states]
                            )
                            cat_vs = _concat_state_token_axis(
                                [vs for _, vs in run_states]
                            )
                            tq = TurboQuantKVCache(bits=tq_bits, seed=tq_seed)
                            tq.keys = cat_ks
                            tq.values = cat_vs
                            tq.offset = _state_length(cat_ks)
                            _rebuild_codecs(tq, cat_ks, cat_vs)
                            reconstructed_caches.append(tq)
                            continue

                        key_parts = []
                        value_parts = []
                        dense_dtype = None
                        for group in groups:
                            if group[0] == "plain":
                                key_parts.append(group[1])
                                value_parts.append(group[2])
                                if dense_dtype is None:
                                    dense_dtype = group[1].dtype
                                continue
                            tq_bits, tq_seed = group[2]
                            dq = self._dequantize_tq_run(
                                group[1], tq_bits, tq_seed, layer_idx
                            )
                            if dq is None:
                                return None
                            key_parts.append(dq[0])
                            value_parts.append(dq[1])

                        keys = mx.concatenate(key_parts, axis=2)
                        values = mx.concatenate(value_parts, axis=2)
                        if dense_dtype is not None and keys.dtype != dense_dtype:
                            keys = keys.astype(dense_dtype)
                            values = values.astype(dense_dtype)
                        cache = KVCache()
                        cache.keys = keys
                        cache.values = values
                        cache.offset = keys.shape[2]
                        reconstructed_caches.append(cache)
                        healed_tq_layers += 1
                        if healed_tq_layers == 1:
                            n_tq = sum(1 for g in groups if g[0] == "tq")
                            n_plain = len(groups) - n_tq
                            logger.info(
                                f"Healed mixed-format TQ chain: layer "
                                f"{layer_idx} has {n_tq} quantized run(s) and "
                                f"{n_plain} dense block group(s); restoring "
                                f"affected layers as dense KVCache "
                                f"(dtype={keys.dtype})"
                            )
                    except Exception as e:
                        logger.error(
                            f"TQ layer {layer_idx}: reconstruction failed: {e}"
                        )
                        return None
                    continue

                # === Generic N-tuple sliceable cache: concatenate block slices ===
                last_block_layer_data = all_block_data[-1][layer_idx]
                if (
                    handler.supports_block_slicing
                    and isinstance(last_block_layer_data, tuple)
                    and len(last_block_layer_data) >= 3
                    and last_block_layer_data[0] == "__nstate__"
                ):
                    marker_class = last_block_layer_data[1] or cache_type_name
                    marker_handler = CacheTypeRegistry.get_handler_by_class_name(
                        marker_class
                    )
                    axis_info = marker_handler.get_state_axis_info()
                    layer_states = []
                    for block_data in all_block_data:
                        if layer_idx >= len(block_data):
                            logger.debug(
                                f"Layer {layer_idx}: missing block data for "
                                f"{marker_class}"
                            )
                            return None
                        block_layer_data = block_data[layer_idx]
                        if (
                            not isinstance(block_layer_data, tuple)
                            or len(block_layer_data) < 3
                            or block_layer_data[0] != "__nstate__"
                        ):
                            logger.debug(
                                f"Layer {layer_idx}: expected N-tuple block data "
                                f"for {marker_class}"
                            )
                            return None
                        elements = tuple(block_layer_data[2])
                        state_dict = {
                            "states": elements,
                            "cache_type": marker_class,
                        }
                        for info, elem in zip(axis_info, elements):
                            state_dict[info.name] = elem
                        layer_states.append(state_dict)

                    concat_state = marker_handler.concatenate_states(layer_states)
                    cache = marker_handler.reconstruct_cache(concat_state, None)
                    if cache is None:
                        logger.error(
                            f"Layer {layer_idx}: failed to reconstruct {marker_class}"
                        )
                        return None
                    reconstructed_caches.append(cache)
                    continue

                # === Generic N-tuple non-sliceable cache: use latest boundary ===
                if (
                    isinstance(last_block_layer_data, tuple)
                    and len(last_block_layer_data) >= 3
                    and last_block_layer_data[0] == "__nstate__"
                ):
                    marker_class = last_block_layer_data[1] or cache_type_name
                    elements = last_block_layer_data[2]
                    marker_handler = CacheTypeRegistry.get_handler_by_class_name(
                        marker_class
                    )
                    meta_state = None
                    if last_block_meta_states and layer_idx < len(
                        last_block_meta_states
                    ):
                        meta_state = last_block_meta_states[layer_idx]
                    cache = marker_handler.deserialize_state(
                        tuple(elements), meta_state
                    )
                    if cache is None:
                        logger.error(
                            f"Layer {layer_idx}: failed to reconstruct {marker_class}"
                        )
                        return None
                    reconstructed_caches.append(cache)
                    continue

                # Collect layer data from all blocks
                layer_states = []
                for block_data in all_block_data:
                    if layer_idx < len(block_data):
                        keys_slice, values_slice = block_data[layer_idx]
                        if keys_slice is not None and values_slice is not None:
                            layer_states.append(
                                {
                                    "keys": keys_slice,
                                    "values": values_slice,
                                }
                            )

                if not layer_states:
                    logger.debug(
                        f"Layer {layer_idx} has no data, cannot reconstruct cache"
                    )
                    return None

                # Get meta_state for this layer based on cache type
                meta_state = None
                if not handler.supports_block_slicing:
                    # Non-sliceable caches (RotatingKVCache, ArraysCache): use LAST block's meta_state
                    # because we use the last block's data (layer_states[-1])
                    if last_block_meta_states and layer_idx < len(
                        last_block_meta_states
                    ):
                        meta_state = last_block_meta_states[layer_idx]
                else:
                    # Sliceable caches (KVCache): first block's meta_state is fine
                    if first_block_meta_states and layer_idx < len(
                        first_block_meta_states
                    ):
                        meta_state = first_block_meta_states[layer_idx]

                # Reconstruct using appropriate handler
                if handler.supports_block_slicing:
                    # Standard concatenation for KVCache
                    concat_state = handler.concatenate_states(layer_states)
                    cache = handler.reconstruct_cache(concat_state, meta_state)
                else:
                    # Non-sliceable cache: use latest state
                    # States were stored as full state, use last one
                    latest_keys = layer_states[-1].get("keys")
                    latest_values = layer_states[-1].get("values")

                    if CacheTypeRegistry.is_rotating_family(cache_type_name):
                        # RotatingKVCache: strict last-block restore.
                        # If the last matched block is a placeholder, we only
                        # had a partial prefix hit and must reject.
                        if hasattr(latest_keys, "shape") and latest_keys.shape == (1,):
                            logger.info(
                                f"RotatingKVCache layer {layer_idx}: partial prefix "
                                f"match detected (placeholder in last matched "
                                f"block). Rejecting cache to prevent stale "
                                f"sliding-window state."
                            )
                            return None

                        latest_state = {
                            "keys": latest_keys,
                            "values": latest_values,
                            "meta_state": meta_state,
                        }
                        cache = handler.reconstruct_cache(latest_state, meta_state)
                    else:
                        # ArraysCache/MambaCache: detect placeholder from
                        # last-block-only storage. If the last matched block
                        # has placeholder shape (1,), this is a partial prefix
                        # match — the real state lives in a later block that
                        # was not matched. We must reject the entire cache
                        # because GDN recurrent state cannot be partially
                        # reconstructed.
                        if hasattr(latest_keys, "shape") and latest_keys.shape == (1,):
                            logger.info(
                                f"ArraysCache layer {layer_idx}: partial prefix "
                                f"match detected (placeholder in last matched "
                                f"block). Rejecting cache to prevent stale GDN "
                                f"state. Request will reprocess from scratch."
                            )
                            return None

                        # Exact match: last block has full state
                        latest_state = {
                            "states": [latest_keys, latest_values],
                        }
                        # Pass token_count for proper SizedArraysCache wrapping
                        cache = handler.reconstruct_cache(
                            latest_state,
                            meta_state,
                            token_count=valid_token_count,
                        )

                if cache is None:
                    # Fallback to simple KVCache reconstruction
                    cache = self._fallback_reconstruct_layer(
                        layer_states, cache_type_name
                    )

                if cache is None:
                    logger.debug(
                        f"Layer {layer_idx}: failed to reconstruct {cache_type_name}"
                    )
                    return None

                reconstructed_caches.append(cache)

            if not reconstructed_caches:
                return None

            # Verify all layers were reconstructed
            if len(reconstructed_caches) != num_layers:
                logger.warning(
                    f"Incomplete cache reconstruction: got {len(reconstructed_caches)} "
                    f"layers, expected {num_layers}"
                )
                return None

            # Verify KVCache offset consistency across KVCache-typed layers.
            # All KVCache layers must have the same offset (they process
            # the same tokens). A mismatch causes broadcast_shapes errors
            # when the model creates a single attention mask from one layer
            # and applies it to all attention layers.
            # NOTE: only check layers explicitly typed as 'KVCache'.
            # RotatingKVCache also has 'offset' but its meaning differs
            # (total tokens ever processed, not buffer size), so mixing
            # them would produce false positives.
            if layer_cache_types:
                kv_offsets = set()
                for idx, c in enumerate(reconstructed_caches):
                    if (
                        idx < len(layer_cache_types)
                        and layer_cache_types[idx] == "KVCache"
                        and hasattr(c, "offset")
                        and isinstance(getattr(c, "offset", None), int)
                    ):
                        kv_offsets.add(c.offset)
                if len(kv_offsets) > 1:
                    logger.warning(
                        f"KVCache offset inconsistency after reconstruction: "
                        f"{kv_offsets}. Rejecting cache to prevent "
                        f"broadcast_shapes errors."
                    )
                    return None

            logger.debug(
                f"Reconstructed cache from tiered cache: {len(reconstructed_caches)} layers, "
                f"{block_table.num_tokens} tokens from {len(block_table.block_ids)} blocks"
            )

            return reconstructed_caches

        except Exception as e:
            logger.warning(f"Failed to reconstruct cache: {e}")
            import traceback

            logger.debug(traceback.format_exc())
            return None

    @staticmethod
    def _layer_has_turboquant_payload(
        all_block_data: list[list[Any]],
        layer_idx: int,
    ) -> bool:
        """True if any block's payload for this layer is TurboQuant-tagged.

        Used to route dense-typed (KVCache) layers whose chain later
        accumulated TQ-converted blocks into the payload-driven TurboQuant
        reconstruction path. layer_cache_types comes from the first block
        only, so the chain-level type cannot detect this mixing.
        """
        for block_data in all_block_data:
            if layer_idx >= len(block_data):
                continue
            bd = block_data[layer_idx]
            if (
                isinstance(bd, tuple)
                and len(bd) == 2
                and isinstance(bd[0], str)
                and bd[0] == "__turboquant_v2__"
            ):
                return True
        return False

    def _tq_block_params(
        self,
        all_block_meta_states: list[Any],
        first_block_meta_states: Any,
        block_idx: int,
        layer_idx: int,
    ) -> tuple[float, int] | None:
        """Resolve (bits, seed) for one block's TurboQuant payload.

        Prefers the block's own per-block meta_state (TurboQuantKVCache
        stores ``(offset, bits, seed)``; mixed chains can carry different
        parameters per block), then the first block's chain-level meta, then
        the server-configured TurboQuant KV bit depth (blocks admitted for
        reconstruction already passed the bit-depth eligibility check, and
        seed 0 is the universal codec default). Returns None when none of
        these resolve: rebuilding a codec at a guessed width dequantizes to
        plausible-but-wrong tensors — silent output corruption after a cache
        hit — so the caller must reject the hit and re-prefill instead.
        """
        candidates = []
        if block_idx < len(all_block_meta_states):
            bms = all_block_meta_states[block_idx]
            if bms and layer_idx < len(bms):
                candidates.append(bms[layer_idx])
        if first_block_meta_states and layer_idx < len(first_block_meta_states):
            candidates.append(first_block_meta_states[layer_idx])
        for ms in candidates:
            if isinstance(ms, (list, tuple)) and len(ms) >= 3:
                try:
                    return float(ms[1]), int(ms[2])
                except (TypeError, ValueError):
                    continue
        expected_bits = getattr(
            self.paged_ssd_cache, "_expected_turboquant_kv_bits", None
        )
        if isinstance(expected_bits, (int, float)):
            from mlx_vlm.turboquant import DEFAULT_TURBOQUANT_SEED

            return float(expected_bits), int(DEFAULT_TURBOQUANT_SEED)
        return None

    def _dequantize_tq_run(
        self,
        run_states: list[tuple[Any, Any]],
        bits: float,
        seed: int,
        layer_idx: int,
    ) -> tuple[Any, Any] | None:
        """Dequantize a run of consecutive TurboQuant block slices to dense KV.

        Concatenates the quantized (ks, vs) states in block order and
        dequantizes once, with codecs rebuilt deterministically from
        (bits, seed). TurboQuant states are per-token (no cross-token
        coupling), so dequantizing a run equals dequantizing its blocks
        individually and concatenating — which makes payload-driven
        per-block reconstruction valid.

        Returns:
            (keys, values) dense float32 arrays, or None when the run cannot
            be decoded (caller rejects the cache hit).
        """
        from mlx_vlm.turboquant import TurboQuantKVCache

        from ..turboquant_kv import _concat_state_token_axis, _rebuild_codecs

        try:
            cat_ks = _concat_state_token_axis([ks for ks, _ in run_states])
            cat_vs = _concat_state_token_axis([vs for _, vs in run_states])
            tq = TurboQuantKVCache(bits=bits, seed=seed)
            _rebuild_codecs(tq, cat_ks, cat_vs)
            return tq.dequantize(cat_ks, cat_vs)
        except Exception as e:
            logger.error(
                f"TQ layer {layer_idx}: failed to dequantize "
                f"{len(run_states)} block state(s) "
                f"(bits={bits}, seed={seed}): {e}"
            )
            return None

    def _fallback_reconstruct_layer(
        self,
        layer_states: list[dict[str, Any]],
        cache_type_name: str,
    ) -> Any | None:
        """
        Fallback layer reconstruction when handler fails.

        Args:
            layer_states: List of state dicts with 'keys' and 'values'
            cache_type_name: Name of the cache type

        Returns:
            Reconstructed cache object or None
        """
        # The fallback rebuilds a plain KVCache, which is the wrong cache
        # class for every other type: a rotating buffer or recurrent state
        # restored as KVCache carries wrong positions and merge-unsafe
        # state. When a non-KVCache handler declares failure, reject the
        # cached prefix instead of papering over it with a guessed rebuild.
        if cache_type_name != "KVCache":
            logger.warning(
                f"Handler reconstruction failed for {cache_type_name}; "
                f"rejecting the cached prefix instead of rebuilding it as "
                f"a plain KVCache."
            )
            return None
        try:
            # Collect keys and values
            layer_keys = [s["keys"] for s in layer_states if s.get("keys") is not None]
            layer_values = [
                s["values"] for s in layer_states if s.get("values") is not None
            ]

            if not layer_keys or not layer_values:
                return None

            # Try to concatenate (works for 4D KV caches)
            try:
                concat_keys = mx.concatenate(layer_keys, axis=2)
                concat_values = mx.concatenate(layer_values, axis=2)
            except Exception:
                # If concatenation fails, might be 3D tensors
                try:
                    concat_keys = mx.concatenate(layer_keys, axis=1)
                    concat_values = mx.concatenate(layer_values, axis=1)
                except Exception:
                    # Last resort: use single state
                    concat_keys = layer_keys[-1]
                    concat_values = layer_values[-1]

            # Create appropriate cache object
            try:
                from mlx_lm.models.cache import KVCache

                cache = KVCache()
                cache.keys = concat_keys
                cache.values = concat_values
                if len(concat_keys.shape) >= 3:
                    cache.offset = (
                        concat_keys.shape[2]
                        if len(concat_keys.shape) == 4
                        else concat_keys.shape[1]
                    )
                else:
                    cache.offset = 0
                return cache
            except ImportError:
                # Simple fallback
                class SimpleKVCache:
                    def __init__(self, keys, values):
                        self.keys = keys
                        self.values = values
                        self.offset = keys.shape[2] if len(keys.shape) >= 3 else 0

                    @property
                    def state(self):
                        return (self.keys, self.values)

                return SimpleKVCache(concat_keys, concat_values)

        except Exception as e:
            logger.debug(f"Fallback reconstruction failed: {e}")
            return None

    def _find_kv_shape_ref(
        self,
        all_block_data: list[list[tuple[Any, Any]]],
        layer_cache_types: list[str] | None = None,
    ) -> tuple[int, int] | None:
        """Find (kv_heads, head_dim) from a KVCache layer's stored data.

        Used to create zero-length RotatingKVCache tensors with the correct shape.

        Args:
            all_block_data: All loaded block data
            layer_cache_types: Per-layer cache type names

        Returns:
            (kv_heads, head_dim) tuple, or None if not found
        """
        if not all_block_data:
            return None

        for layer_idx, layer_data in enumerate(all_block_data[0]):
            # Skip non-KVCache layers
            if layer_cache_types and layer_idx < len(layer_cache_types):
                if layer_cache_types[layer_idx] != "KVCache":
                    continue
            # Guard against non-tuple formats (CacheList stores List[Tuple])
            if not isinstance(layer_data, tuple) or len(layer_data) != 2:
                continue
            keys, _ = layer_data
            if hasattr(keys, "shape") and len(keys.shape) == 4:
                return (keys.shape[1], keys.shape[3])

        return None

    def _create_empty_rotating_cache(
        self,
        meta_state: tuple | None = None,
        kvcache_offset: int = 0,
        kv_shape_ref: tuple[int, int] | None = None,
    ) -> Any | None:
        """
        Create an empty RotatingKVCache for partial prefix restore.

        Creates a RotatingKVCache with zero-length keys/values (not None) and
        offset matching the KVCache layers. This ensures:
        1. mlx-lm's empty() returns False → Continuation mode (not Fresh Start)
        2. Position IDs (RoPE) are correct for all layers
        3. The merge creates a zero-length buffer (not zero-filled) so that
           no phantom attention positions exist during window padding reprocessing

        Uses PrefillReadyRotatingKVCache (clamped size() by buffer length)
        so BatchRotatingKVCache.merge() never reads beyond the actual buffer
        and never sees zero-padded positions as valid attention keys.

        Args:
            meta_state: RotatingKVCache meta_state tuple (keep, max_size, offset, _idx).
            kvcache_offset: Offset to match KVCache layers (= restored token count).
            kv_shape_ref: (kv_heads, head_dim) from a KVCache layer for tensor shape.

        Returns:
            RotatingKVCache with zero-length keys/values, or None on failure.
        """
        if meta_state and len(meta_state) >= 2:
            keep = int(meta_state[0])
            max_size = int(meta_state[1])
        else:
            logger.warning(
                "Cannot create empty RotatingKVCache: meta_state missing or incomplete"
            )
            return None

        cache = PrefillReadyRotatingKVCache(max_size=max_size, keep=keep)
        cache.offset = kvcache_offset

        # Set zero-length keys/values so empty() returns False.
        # This prevents mlx-lm from entering Fresh Start mode which
        # would discard all cached KVCache data.
        if kv_shape_ref and HAS_MLX:
            kv_heads, head_dim = kv_shape_ref
            cache.keys = mx.zeros((1, kv_heads, 0, head_dim))
            cache.values = mx.zeros((1, kv_heads, 0, head_dim))
            cache._idx = 0
            logger.debug(
                f"Created empty RotatingKVCache: max_size={max_size}, keep={keep}, "
                f"offset={kvcache_offset}, kv_heads={kv_heads}, head_dim={head_dim}"
            )
        else:
            logger.debug(
                f"Created empty RotatingKVCache: max_size={max_size}, keep={keep} "
                f"(no shape ref, keys=None)"
            )

        return cache

    def _validate_block_cache_data(
        self,
        cache_data: list[tuple[Any, Any]],
        layer_cache_types: list[str] | None = None,
    ) -> bool:
        """
        Validate that block's cache_data has valid data for all layers.

        A block's cache_data is a list of (keys, values) tuples, one per layer.
        This validates that:
        1. cache_data is not empty
        2. Each layer has non-None keys and values
        3. Each layer has consistent shapes (for sliceable cache types)

        Args:
            cache_data: List of (keys, values) tuples from a block
            layer_cache_types: Optional list of cache type names per layer.
                ArraysCache layers are excluded from seq_len consistency check.

        Returns:
            True if valid, False otherwise
        """
        if not cache_data:
            return False

        # Cache types that support per-block KV slicing. When cache type
        # metadata is present, do not let non-KV hybrid states define the
        # expected seq_len for KVCache layers.
        sliceable_types = {
            "KVCache",
            "BatchKVCache",
            "TurboQuantKVCache",
            "BatchTurboQuantKVCache",
            "MiniMaxM3KVCache",
        }
        non_sliceable_types = {
            "ArraysCache",
            "RotatingKVCache",
            "BatchRotatingKVCache",
            "CacheList",
            "MiniMaxM3BatchKVCache",
        }

        expected_seq_len = None

        for layer_idx, layer_data in enumerate(cache_data):
            try:
                # Determine cache type first to handle CacheList before tuple unpack
                cache_type = None
                if layer_cache_types and layer_idx < len(layer_cache_types):
                    cache_type = layer_cache_types[layer_idx]

                # CacheList: sub-cache list format, skip standard (keys, values) unpacking
                if cache_type == "CacheList":
                    # CacheList data is either List[Tuple] (last block) or Tuple (placeholder)
                    if isinstance(layer_data, list):
                        continue  # Sub-cache list — valid
                    # Fall through to standard check for placeholder (zeros tuple)

                if not isinstance(layer_data, (list, tuple)) or len(layer_data) < 2:
                    logger.debug(
                        f"Block validation failed: layer {layer_idx} has "
                        f"unsupported data type {type(layer_data).__name__}"
                    )
                    return False

                if (
                    isinstance(layer_data, tuple)
                    and len(layer_data) >= 3
                    and layer_data[0] == "__nstate__"
                ):
                    elements = layer_data[2]
                    if not isinstance(elements, (list, tuple)) or len(elements) < 2:
                        logger.debug(
                            f"Block validation failed: layer {layer_idx} has "
                            "invalid N-tuple state"
                        )
                        return False
                    keys, values = elements[0], elements[1]
                else:
                    keys, values = layer_data[0], layer_data[1]

                # Skip seq_len check for non-sliceable types (e.g., ArraysCache, RotatingKVCache)
                # This includes placeholder entries (1D tensors from non-last blocks)
                # used by the last-block-only RotatingKVCache storage strategy
                if cache_type in non_sliceable_types:
                    continue
                if layer_cache_types and cache_type not in sliceable_types:
                    continue

                # Check for None after non-sliceable N-tuple caches have been
                # accepted. Pooling-style states can legitimately store None
                # in their first elements.
                if keys is None or values is None:
                    logger.debug(
                        f"Block validation failed: layer {layer_idx} has None keys/values"
                    )
                    return False

                # Check shape consistency for sliceable types (KVCache, RotatingKVCache)
                if hasattr(keys, "shape") and len(keys.shape) >= 3:
                    seq_len = keys.shape[2]
                    if expected_seq_len is None:
                        expected_seq_len = seq_len
                    elif seq_len != expected_seq_len:
                        logger.debug(
                            f"Block validation failed: layer {layer_idx} has "
                            f"seq_len {seq_len}, expected {expected_seq_len}"
                        )
                        return False
            except (TypeError, ValueError) as e:
                logger.debug(f"Block validation failed: layer {layer_idx} error: {e}")
                return False

        return True

    def _find_best_prefix_match(
        self,
        tokens: list[int],
        extra_keys: tuple[Any, ...] | None = None,
    ) -> tuple[int, tuple[int, ...], int, list[bytes]] | None:
        """Find best matching prefix in the index.

        Returns:
            Tuple of (prefix_len, block_ids, num_blocks, chain_hashes) where
            chain_hashes holds the per-block chain hash for each matched
            block, or None when nothing matched. The hashes let the caller
            re-validate that the indexed blocks still hold the content they
            were indexed for before reusing them.
        """
        best_match = None
        best_len = 0

        parent_hash = b""
        prefix_len = 0
        num_blocks = 0
        chain_hashes: list[bytes] = []

        for start in range(0, len(tokens), self.block_size):
            end = min(start + self.block_size, len(tokens))
            block_tokens = tokens[start:end]
            if not block_tokens:
                break

            parent_hash = compute_block_hash(
                parent_hash,
                block_tokens,
                extra_keys=extra_keys,
                model_name=self.paged_cache.model_name,
            )
            prefix_len += len(block_tokens)
            num_blocks += 1
            chain_hashes.append(parent_hash)

            entry = self._prefix_index.get(parent_hash)
            if entry and entry[0] == prefix_len and prefix_len > best_len:
                best_match = entry
                best_len = prefix_len

        if best_match is None:
            return None
        return (*best_match, chain_hashes[: best_match[2]])

    def _update_prefix_index(
        self,
