"""
OpenAI-compatible local inference server with multi-engine auto-benchmark.

Auto-detects available engines, benchmarks each, and picks the fastest:
  1. mlx-lm (Apple Silicon — quantized HuggingFace models)
  2. llama-cpp-python (cross-platform — GGUF models)
  3. MNN-LLM (Metal/Vulkan/OpenCL/CUDA — via engine plugin)
  4. ONNX Runtime (CPU/CUDA/DirectML/TensorRT — via engine plugin)
  5. Echo (fallback — for testing the API layer)

Usage::

    from octomil.serve import run_server
    run_server("gemma-2b", port=8080)

    # Override engine selection:
    run_server("gemma-2b", port=8080, engine="llama.cpp")

    # Then:
    # curl localhost:8080/v1/chat/completions \\
    #   -d '{"model":"gemma-2b","messages":[{"role":"user","content":"Hi"}]}'
"""

from __future__ import annotations

import json
import logging
import os
import time
import uuid
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, AsyncIterator, Optional, Union

if TYPE_CHECKING:
    from .early_exit import EarlyExitConfig, EarlyExitMonitor
    from .telemetry import TelemetryReporter

from pydantic import BaseModel, Field

from .models.catalog import CATALOG as _UNIFIED_CATALOG
from .models.resolver import ModelResolutionError as _NewResolutionError
from .models.resolver import resolve as _resolve_new

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Pydantic models for OpenAI-compatible API
# ---------------------------------------------------------------------------


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatCompletionBody(BaseModel):
    model: str = ""
    messages: list[ChatMessage] = Field(default_factory=list)
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 1.0
    stream: bool = False
    response_format: Optional[dict[str, Any]] = None
    grammar: Optional[str] = None


# ---------------------------------------------------------------------------
# Model catalog — backwards-compatible dicts derived from unified catalog
# ---------------------------------------------------------------------------

_MLX_MODELS: dict[str, str] = {}
_GGUF_MODELS: dict[str, tuple[str, str]] = {}

for _name, _entry in _UNIFIED_CATALOG.items():
    _default_variant = _entry.variants.get(_entry.default_quant)
    if _default_variant is None:
        continue
    if _default_variant.mlx:
        _MLX_MODELS[_name] = _default_variant.mlx
    if _default_variant.gguf:
        _GGUF_MODELS[_name] = (
            _default_variant.gguf.repo,
            _default_variant.gguf.filename,
        )


def resolve_model_name(name: str, backend: str) -> str:
    """Resolve a short model name (with optional :variant) to a HuggingFace repo ID.

    Uses the unified model catalog for structured resolution. Supports
    Ollama-style ``model:variant`` syntax (e.g. ``gemma-3b:4bit``).

    Full repo paths (containing ``/``) and local file paths pass through
    unchanged.
    """
    # Local file path
    if name.endswith((".gguf", ".pte", ".mnn")):
        return name

    # Full repo path — pass through
    if "/" in name:
        return name

    # Map backend names to engine names for the resolver
    engine_map = {"mlx": "mlx-lm", "gguf": "llama.cpp"}
    engine = engine_map.get(backend, backend)

    try:
        resolved = _resolve_new(name, engine=engine)
    except _NewResolutionError as exc:
        raise ValueError(str(exc)) from exc

    if backend == "mlx":
        if resolved.mlx_repo:
            return resolved.mlx_repo
        if resolved.hf_repo:
            return resolved.hf_repo
        raise ValueError(f"No MLX source found for '{name}'. " f"Pass a full HuggingFace repo ID (e.g. REDACTED)")

    if backend == "gguf":
        # For GGUF, return the short name — LlamaCppBackend resolves via _GGUF_MODELS
        family = resolved.family
        if family and family in _GGUF_MODELS:
            return family
        # If the resolver found a GGUF artifact, return the family name
        if resolved.is_gguf and family:
            return family
        raise ValueError(
            f"No GGUF source found for '{name}'. " f"Pass a path to a local .gguf file or a HuggingFace repo ID."
        )

    return name


# ---------------------------------------------------------------------------
# Backend abstraction
# ---------------------------------------------------------------------------


@dataclass
class GenerationRequest:
    model: str
    messages: list[dict[str, str]]
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 1.0
    stream: bool = False
    grammar: Optional[str] = None
    json_mode: bool = False


@dataclass
class GenerationChunk:
    text: str
    token_count: int = 0
    tokens_per_second: float = 0.0
    finish_reason: Optional[str] = None


@dataclass
class InferenceMetrics:
    """Metrics for a single inference request."""

    ttfc_ms: float = 0.0
    prompt_tokens: int = 0
    total_tokens: int = 0
    tokens_per_second: float = 0.0
    total_duration_ms: float = 0.0
    cache_hit: bool = False
    attention_backend: str = "standard"
    early_exit_tokens: int = 0
    avg_layers_used: float = 0.0


class InferenceBackend:
    """Base class for inference backends."""

    name: str = "base"
    attention_backend: str = "standard"

    def load_model(self, model_name: str) -> None:
        raise NotImplementedError

    def generate(self, request: GenerationRequest) -> tuple[str, InferenceMetrics]:
        raise NotImplementedError

    async def generate_stream(
        self,
        request: GenerationRequest,
    ) -> AsyncIterator[GenerationChunk]:
        raise NotImplementedError
        yield  # pragma: no cover — makes this an async generator

    def list_models(self) -> list[str]:
        raise NotImplementedError


_DRAFT_MODEL_MAP: dict[str, str] = {
    "phi-4": "REDACTED",
    "llama-8b": "REDACTED",
    "llama-3b": "REDACTED",
    "qwen-7b": "REDACTED",
    "qwen-3b": "REDACTED",
    "gemma-4b": "REDACTED",
    "gemma-12b": "REDACTED",
    "gemma-27b": "REDACTED",
}


class MLXBackend(InferenceBackend):
    """Apple Silicon backend using mlx-lm.

    Loads quantized models from HuggingFace (auto-downloads + caches).
    Uses the tokenizer's built-in chat template for proper formatting.
    Supports KV cache persistence for prefix reuse across requests.

    MLX uses Metal fused attention automatically on Apple Silicon — no
    explicit configuration needed.  Reported as ``metal_fused``.
    """

    name = "mlx-lm"
    attention_backend = "metal_fused"

    def __init__(
        self,
        cache_size_mb: int = 2048,
        cache_enabled: bool = True,
    ) -> None:
        self._model: Any = None
        self._tokenizer: Any = None
        self._draft_model: Any = None
        self._model_name: str = ""
        self._repo_id: str = ""
        self._cache_enabled = cache_enabled
        # Multi-entry KV cache pool — LRU eviction by both entry count and size.
        from .cache import KVCacheManager

        self._cache_mgr = KVCacheManager(max_cache_size_mb=cache_size_mb, max_entries=4)
        self._cache_hits: int = 0
        self._cache_misses: int = 0

    def load_model(self, model_name: str) -> None:
        import mlx_lm  # type: ignore[import-untyped]

        self._model_name = model_name
        self._repo_id = resolve_model_name(model_name, "mlx")

        logger.info("Loading %s (%s) with mlx-lm...", model_name, self._repo_id)
        self._model, self._tokenizer, *_ = mlx_lm.load(self._repo_id)
        logger.info("Model loaded: %s", self._repo_id)

        # Set Metal buffer cache limit to prevent unbounded growth
        # during long inference sessions.  2GB is enough for ~14B models.
        try:
            import mlx.core as mx  # type: ignore[import-untyped]

            mx.set_cache_limit(2 * 1024 * 1024 * 1024)  # 2 GB
        except Exception:
            pass

        # Load draft model for speculative decoding (if available)
        draft_repo = _DRAFT_MODEL_MAP.get(model_name)
        if draft_repo:
            logger.info("Loading draft model for speculative decoding: %s", draft_repo)
            try:
                self._draft_model, *_ = mlx_lm.load(draft_repo)
            except Exception:
                logger.debug("Draft model load failed (non-fatal)", exc_info=True)

        # Collect special token strings for output filtering
        self._stop_strings: set[str] = set()
        if hasattr(self._tokenizer, "eos_token") and self._tokenizer.eos_token:
            self._stop_strings.add(self._tokenizer.eos_token)
        # Common chat model turn markers
        for marker in ("<end_of_turn>", "<|eot_id|>", "<|im_end|>", "</s>"):
            self._stop_strings.add(marker)

    def _apply_chat_template(self, messages: list[dict[str, str]]) -> str:
        """Format messages using the model's chat template."""
        try:
            result: str = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            return result
        except Exception:
            # Fallback for models without a chat template
            parts: list[str] = []
            for msg in messages:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                parts.append(f"{role}: {content}")
            parts.append("assistant:")
            return "\n".join(parts)

    def _make_sampler(self, temperature: float, top_p: float) -> Any:
        """Create an mlx-lm sampler with the given temperature and top_p."""
        from mlx_lm.sample_utils import make_sampler  # type: ignore[import-untyped]

        return make_sampler(temp=temperature, top_p=top_p)

    def _is_stop_token(self, text: str) -> bool:
        """Check if the token text is a stop/EOS marker."""
        stripped = text.strip()
        return stripped in self._stop_strings

    def _fetch_or_create_cache(self, prompt_token_ids: list[int]) -> tuple[list[Any], list[int], bool]:
        """Fetch a reusable prompt cache or create a fresh one.

        Returns (cache, remaining_tokens, cache_hit).  Looks up the longest
        matching prefix in the ``KVCacheManager`` pool, trims the MLX KV
        tensors to the common prefix, and returns only the remaining tokens
        so ``stream_generate`` only processes the delta.
        """
        from mlx_lm.models.cache import (  # type: ignore[import-untyped]
            can_trim_prompt_cache,
            make_prompt_cache,
            trim_prompt_cache,
        )

        if self._cache_enabled:
            hit = self._cache_mgr.get_cached_prefix(prompt_token_ids)
            if hit is not None:
                cache = hit.kv_state
                common_len = hit.prefix_length

                # Trim cache to the reusable prefix.
                # For exact match (common_len == len(prompt)), trim to
                # common_len - 1 and re-feed the last token, because
                # stream_generate requires a non-empty prompt.
                trim_target = common_len
                if common_len == len(prompt_token_ids):
                    trim_target = common_len - 1

                total_cached = cache[0].offset
                num_to_trim = total_cached - trim_target
                if num_to_trim > 0:
                    if can_trim_prompt_cache(cache):
                        trim_prompt_cache(cache, num_to_trim)
                    else:
                        cache = make_prompt_cache(self._model)
                        self._cache_misses += 1
                        return cache, prompt_token_ids, False

                remaining = prompt_token_ids[trim_target:]
                logger.debug(
                    "MLX cache hit: reusing %d/%d prompt tokens (%d to process)",
                    trim_target,
                    len(prompt_token_ids),
                    len(remaining),
                )
                self._cache_hits += 1
                return cache, remaining, True

        # No usable cache — create fresh
        self._cache_misses += 1
        return make_prompt_cache(self._model), prompt_token_ids, False

    def _store_cache(self, prompt_token_ids: list[int], cache: list[Any]) -> None:
        """Store the prompt cache in the pool for potential reuse."""
        if self._cache_enabled:
            self._cache_mgr.store_prefix(prompt_token_ids, cache)

    def generate(self, request: GenerationRequest) -> tuple[str, InferenceMetrics]:
        import mlx_lm  # type: ignore[import-untyped]

        prompt = self._apply_chat_template(request.messages)
        prompt_token_ids = self._tokenizer.encode(prompt)
        prompt_tokens = len(prompt_token_ids)
        sampler = self._make_sampler(request.temperature, request.top_p)
        start = time.monotonic()
        first_token_time: Optional[float] = None
        tokens: list[str] = []
        final_tps: float = 0.0

        # Fetch reusable prompt cache or create fresh
        cache, remaining_tokens, cache_hit = self._fetch_or_create_cache(prompt_token_ids)

        extra_kwargs: dict[str, Any] = {
            "prompt_cache": cache,
            "prefill_step_size": 4096,
        }
        if self._draft_model is not None:
            extra_kwargs["draft_model"] = self._draft_model

        for response in mlx_lm.stream_generate(
            self._model,
            self._tokenizer,
            prompt=remaining_tokens,
            max_tokens=request.max_tokens,
            sampler=sampler,
            **extra_kwargs,
        ):
            if first_token_time is None:
                first_token_time = time.monotonic()
            final_tps = response.generation_tps
            if response.finish_reason or self._is_stop_token(response.text):
                break
            tokens.append(response.text)

        # Store cache for next request (cache was updated in-place)
        self._store_cache(prompt_token_ids, cache)

        elapsed = time.monotonic() - start
        ttfc = ((first_token_time or start) - start) * 1000
        text = "".join(tokens)

        return (
            text,
            InferenceMetrics(
                ttfc_ms=ttfc,
                prompt_tokens=prompt_tokens,
                total_tokens=len(tokens),
                tokens_per_second=final_tps,
                total_duration_ms=elapsed * 1000,
                cache_hit=cache_hit,
                attention_backend="metal_fused",  # MLX uses Metal fused attention automatically
            ),
        )

    async def generate_stream(
        self,
        request: GenerationRequest,
    ) -> AsyncIterator[GenerationChunk]:
        import asyncio
        import threading

        import mlx_lm  # type: ignore[import-untyped]

        prompt = self._apply_chat_template(request.messages)
        prompt_token_ids = self._tokenizer.encode(prompt)
        sampler = self._make_sampler(request.temperature, request.top_p)

        # Fetch reusable prompt cache or create fresh
        cache, remaining_tokens, _cache_hit = self._fetch_or_create_cache(prompt_token_ids)

        extra_kwargs: dict[str, Any] = {
            "prompt_cache": cache,
            "prefill_step_size": 4096,
        }
        if self._draft_model is not None:
            extra_kwargs["draft_model"] = self._draft_model

        # mlx_lm.stream_generate is a sync generator — push chunks from
        # a background thread into an asyncio.Queue for non-blocking consumption.
        # asyncio.Queue is NOT thread-safe, so use call_soon_threadsafe for puts.
        queue: asyncio.Queue[Optional[Any]] = asyncio.Queue()
        cancelled = threading.Event()
        done = threading.Event()
        loop = asyncio.get_event_loop()

        def _produce() -> None:
            try:
                for response in mlx_lm.stream_generate(
                    self._model,
                    self._tokenizer,
                    prompt=remaining_tokens,
                    max_tokens=request.max_tokens,
                    sampler=sampler,
                    **extra_kwargs,
                ):
                    if cancelled.is_set():
                        break
                    loop.call_soon_threadsafe(queue.put_nowait, response)
            finally:
                loop.call_soon_threadsafe(queue.put_nowait, None)  # sentinel
                done.set()

        threading.Thread(target=_produce, daemon=True).start()

        try:
            while True:
                response = await queue.get()
                if response is None:
                    break
                if response.finish_reason or self._is_stop_token(response.text):
                    yield GenerationChunk(
                        text="",
                        finish_reason="stop",
                    )
                    break
                yield GenerationChunk(
                    text=response.text,
                    token_count=response.generation_tokens,
                    tokens_per_second=response.generation_tps,
                )
        finally:
            cancelled.set()
            # Wait for producer to fully exit before reusing the model/cache.
            # Without this, Metal command buffers from the dying producer
            # collide with the next request's Metal operations → SIGABRT.
            done.wait(timeout=5.0)
            # Store cache for next request. The done.wait() above guarantees
            # the producer has fully stopped and all Metal ops are complete,
            # so the cache is in a consistent (if partially filled) state.
            self._store_cache(prompt_token_ids, cache)

    @property
    def kv_cache(self) -> Any:
        """Expose the KVCacheManager for the stats endpoint."""
        return self._cache_mgr

    def list_models(self) -> list[str]:
        return [self._model_name] if self._model_name else []


class LlamaCppBackend(InferenceBackend):
    """Cross-platform backend using llama-cpp-python.

    Loads GGUF models from HuggingFace via from_pretrained (auto-downloads).
    Chat templates are handled by llama.cpp internally.
    Supports built-in LlamaCache for automatic KV prefix reuse.
    Flash attention is enabled by default (``flash_attn=True``) for better
    performance on long sequences.
    """

    name = "llama.cpp"
    attention_backend = "flash_attention"

    def __init__(
        self,
        cache_size_mb: int = 2048,
        cache_enabled: bool = True,
    ) -> None:
        self._llm: Any = None
        self._model_name: str = ""
        self._cache_size_mb = cache_size_mb
        self._cache_enabled = cache_enabled
        self._llama_cache: Any = None

    def _attach_cache(self) -> None:
        """Attach a LlamaCache to the loaded model if caching is enabled."""
        if not self._cache_enabled or self._llm is None:
            return
        try:
            from llama_cpp import LlamaCache  # type: ignore[import-untyped]

            capacity_bytes = self._cache_size_mb * 1024 * 1024
            self._llama_cache = LlamaCache(capacity_bytes=capacity_bytes)
            self._llm.set_cache(self._llama_cache)
            logger.info("LlamaCache enabled: %d MB capacity", self._cache_size_mb)
        except (ImportError, AttributeError, TypeError) as exc:
            logger.debug("LlamaCache not available: %s", exc)

    def load_model(self, model_name: str) -> None:
        from llama_cpp import Llama  # type: ignore[import-untyped]

        self._model_name = model_name

        # Local GGUF file path
        if model_name.endswith(".gguf"):
            logger.info("Loading local GGUF: %s", model_name)
            self._llm = Llama(
                model_path=model_name,
                n_ctx=4096,
                n_batch=256,
                n_gpu_layers=-1,  # offload all layers to GPU
                flash_attn=True,  # fused attention kernels for better perf on long sequences
                verbose=False,
            )
            self._attach_cache()
            return

        # Full HuggingFace repo ID (user/repo)
        if "/" in model_name:
            logger.info("Loading from HuggingFace: %s", model_name)
            self._llm = Llama.from_pretrained(
                repo_id=model_name,
                filename="*Q4_K_M.gguf",
                n_ctx=4096,
                n_batch=256,
                n_gpu_layers=-1,
                flash_attn=True,
                verbose=False,
            )
            self._attach_cache()
            return

        # Try the new resolver for model:variant syntax and catalog lookup
        try:
            resolved = _resolve_new(model_name, engine="llama.cpp")
            if resolved.is_gguf and resolved.filename:
                repo_id = resolved.hf_repo
                filename = resolved.filename
                logger.info(
                    "Loading %s from %s (%s)...",
                    model_name,
                    repo_id,
                    filename,
                )
                self._llm = Llama.from_pretrained(
                    repo_id=repo_id,
                    filename=filename,
                    n_ctx=4096,
                    n_batch=256,
                    n_gpu_layers=-1,
                    flash_attn=True,
                    verbose=False,
                )
                self._attach_cache()
                return
        except _NewResolutionError:
            pass

        # Fallback: short name → legacy catalog lookup
        if model_name not in _GGUF_MODELS:
            raise ValueError(
                f"Unknown model '{model_name}'. "
                f"Available: {', '.join(sorted(_GGUF_MODELS))}\n"
                f"Or pass a local .gguf path or HuggingFace repo ID."
            )

        repo_id, filename = _GGUF_MODELS[model_name]
        logger.info("Loading %s from %s (%s)...", model_name, repo_id, filename)
        self._llm = Llama.from_pretrained(
            repo_id=repo_id,
            filename=filename,
            n_ctx=4096,
            n_batch=256,
            n_gpu_layers=-1,
            flash_attn=True,
            verbose=False,
        )
        self._attach_cache()
        logger.info("Model loaded: %s", model_name)

    def _grammar_arg(self, request: GenerationRequest) -> dict[str, Any]:
        """Build the grammar kwarg for create_chat_completion, if any."""
        if not request.grammar:
            return {}
        try:
            from llama_cpp import LlamaGrammar  # type: ignore[import-untyped]

            return {"grammar": LlamaGrammar.from_string(request.grammar)}
        except Exception as exc:
            logger.warning("Failed to compile GBNF grammar: %s", exc)
            return {}

    def generate(self, request: GenerationRequest) -> tuple[str, InferenceMetrics]:
        start = time.monotonic()
        grammar_kw = self._grammar_arg(request)
        result = self._llm.create_chat_completion(
            messages=request.messages,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            **grammar_kw,
        )
        elapsed = time.monotonic() - start
        text = result["choices"][0]["message"]["content"]
        usage = result.get("usage", {})
        prompt_tokens = usage.get("prompt_tokens", 0)
        completion_tokens = usage.get("completion_tokens", 0)

        return text, InferenceMetrics(
            prompt_tokens=prompt_tokens,
            total_tokens=completion_tokens,
            tokens_per_second=completion_tokens / elapsed if elapsed > 0 else 0,
            total_duration_ms=elapsed * 1000,
            attention_backend="flash_attention",
        )

    async def generate_stream(
        self,
        request: GenerationRequest,
    ) -> AsyncIterator[GenerationChunk]:
        import asyncio

        loop = asyncio.get_event_loop()

        grammar_kw = self._grammar_arg(request)
        # create_chat_completion with stream=True returns a sync iterator
        stream = self._llm.create_chat_completion(
            messages=request.messages,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stream=True,
            **grammar_kw,
        )

        def _next_chunk() -> Any:
            try:
                return next(stream)
            except StopIteration:
                return None

        while True:
            chunk = await loop.run_in_executor(None, _next_chunk)
            if chunk is None:
                break
            delta = chunk["choices"][0].get("delta", {})
            content = delta.get("content", "")
            finish = chunk["choices"][0].get("finish_reason")
            if content or finish:
                yield GenerationChunk(text=content, finish_reason=finish)

    def list_models(self) -> list[str]:
        return [self._model_name] if self._model_name else []


class EchoBackend(InferenceBackend):
    """Fallback backend that echoes input — useful for testing the API layer."""

    name = "echo"

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

    def load_model(self, model_name: str) -> None:
        self._model_name = model_name
        logger.warning(
            "No inference backend available. Using echo backend for '%s'. "
            "Install mlx-lm (Apple Silicon) or llama-cpp-python for real inference: "
            "pip install 'octomil-sdk[mlx]' or pip install 'octomil-sdk[llama]'",
            model_name,
        )

    def generate(self, request: GenerationRequest) -> tuple[str, InferenceMetrics]:
        last_msg = request.messages[-1]["content"] if request.messages else ""
        text = f"[echo:{self._model_name}] {last_msg}"
        return text, InferenceMetrics(total_tokens=len(text.split()))

    async def generate_stream(
        self,
        request: GenerationRequest,
    ) -> AsyncIterator[GenerationChunk]:
        import asyncio

        last_msg = request.messages[-1]["content"] if request.messages else ""
        words = f"[echo:{self._model_name}] {last_msg}".split()
        for i, word in enumerate(words):
            is_last = i == len(words) - 1
            yield GenerationChunk(
                text=word + ("" if is_last else " "),
                finish_reason="stop" if is_last else None,
            )
            await asyncio.sleep(0.02)

    def list_models(self) -> list[str]:
        return [self._model_name] if self._model_name else []


def _detect_backend(
    model_name: str,
    *,
    cache_size_mb: int = 2048,
    cache_enabled: bool = True,
    engine_override: Optional[str] = None,
) -> InferenceBackend:
    """Auto-detect engines, benchmark each, and return the fastest backend.

    Uses the engine registry plugin system. Each registered engine is:
    1. Detected (is the library installed? does it support this model?)
    2. Benchmarked (quick 32-token generation to measure tok/s)
    3. Ranked (highest tok/s wins)

    If engine_override is set, skip benchmarking and use that engine directly.
    """
    from .engines import get_registry

    registry = get_registry()

    backend_kwargs = {
        "cache_size_mb": cache_size_mb,
        "cache_enabled": cache_enabled,
    }

    if engine_override:
        engine = registry.get_engine(engine_override)
        if engine is None:
            available = [e.name for e in registry.engines]
            raise ValueError(f"Unknown engine '{engine_override}'. Available: {', '.join(available)}")
        backend: InferenceBackend = engine.create_backend(model_name, **backend_kwargs)
        return backend

    # Detect all available engines for this model
    detections = registry.detect_all(model_name)
    available_engines = [d.engine for d in detections if d.available]

    for d in detections:
        if d.available:
            logger.info("Engine detected: %s (%s)", d.engine.name, d.info)
        else:
            logger.debug("Engine unavailable: %s", d.engine.name)

    # No real engines → echo fallback
    real_engines = [e for e in available_engines if e.name != "echo"]
    if not real_engines:
        echo = EchoBackend()
        echo.load_model(model_name)
        return echo

    # Benchmark real engines and pick fastest
    ranked = registry.benchmark_all(model_name, n_tokens=32, engines=real_engines)
    best = registry.select_best(ranked)
    if best is None:
        echo = EchoBackend()
        echo.load_model(model_name)
        return echo

    best_backend: InferenceBackend = best.engine.create_backend(model_name, **backend_kwargs)
    return best_backend


def _log_startup_error(model_name: str, exc: Exception) -> None:
    """Print a human-readable startup error instead of a raw traceback."""
    err_type = type(exc).__name__
    err_msg = str(exc)

    # HuggingFace auth / repo errors
    if "RepositoryNotFoundError" in err_type or "401" in err_msg or "403" in err_msg:
        logger.error(
            "Failed to load model '%s': HuggingFace authentication required.\n"
            "  Fix: Run `huggingface-cli login` or set the HF_TOKEN env var.\n"
            "  Get a token at https://huggingface.co/settings/tokens",
            model_name,
        )
    elif "404" in err_msg or "not found" in err_msg.lower():
        logger.error(
            "Failed to load model '%s': model not found on HuggingFace.\n"
            "  Check the model name and try a full repo ID:\n"
            "    octomil serve REDACTED\n"
            "  List available short names:\n"
            "    octomil serve --help",
            model_name,
        )
    elif "Unknown model" in err_msg:
        logger.error("Failed to load model: %s", err_msg)
    else:
        logger.error(
            "Failed to start server for model '%s': %s\n"
            "  If this is a HuggingFace model, try: huggingface-cli login",
            model_name,
            exc,
        )


def _get_cache_manager(backend: InferenceBackend) -> Any:
    """Extract the KVCacheManager from a backend, if available."""
    if isinstance(backend, MLXBackend):
        return backend.kv_cache
    return None


# ---------------------------------------------------------------------------
# FastAPI app factory
# ---------------------------------------------------------------------------


@dataclass
class MoEConfig:
    """Configuration for Mixture of Experts models.

    Controls expert memory management and telemetry behaviour.
    Both llama.cpp and MLX handle MoE natively, so these settings
    primarily control detection, logging, and telemetry — not
    the actual expert routing algorithm.
    """

    enabled: bool = True  # enable MoE-aware features (detection, telemetry)
    expert_memory_limit_mb: int = 0  # 0 = no limit; cap RAM for experts
    log_expert_routing: bool = False  # log per-request expert activation
    offload_inactive: bool = False  # hint to offload inactive experts to disk


@dataclass
class ServerState:
    """Shared mutable state for the serve app."""

    backend: Optional[InferenceBackend] = None
    whisper_backend: Any = None  # _WhisperBackend instance (speech-to-text)
    model_name: str = ""
    engine_name: str = ""
    start_time: float = field(default_factory=time.time)
    request_count: int = 0
    api_key: Optional[str] = None
    api_base: str = "https://api.octomil.com/api/v1"
    default_json_mode: bool = False
    cache_size_mb: int = 2048
    cache_enabled: bool = True
    engine_override: Optional[str] = None
    reporter: Optional["TelemetryReporter"] = None
    max_queue_depth: int = 32
    request_queue: Any = None  # RequestQueue instance
    moe_config: MoEConfig = field(default_factory=MoEConfig)
    is_moe_model: bool = False
    moe_metadata: Any = None  # MoEMetadata from catalog
    compressor: Any = None  # PromptCompressor instance
    early_exit_config: Optional["EarlyExitConfig"] = None
    early_exit_monitor: Optional["EarlyExitMonitor"] = None
    tool_use: bool = False  # pre-load coding agent tool schemas


@dataclass
class MultiModelServerState:
    """Shared mutable state for multi-model serving with routing."""

    backends: dict[str, InferenceBackend] = field(default_factory=dict)
    model_names: list[str] = field(default_factory=list)
    router: Any = None  # QueryRouter instance
    start_time: float = field(default_factory=time.time)
    request_count: int = 0
    routed_counts: dict[str, int] = field(default_factory=dict)
    fallback_counts: int = 0
    api_key: Optional[str] = None
    api_base: str = "https://api.octomil.com/api/v1"
    default_json_mode: bool = False
    cache_size_mb: int = 2048
    cache_enabled: bool = True
    engine_override: Optional[str] = None
    reporter: Optional["TelemetryReporter"] = None
    route_strategy: str = "complexity"
    compressor: Any = None  # PromptCompressor instance


def _resolve_grammar(body: ChatCompletionBody, default_json_mode: bool = False) -> tuple[Optional[str], bool]:
    """Determine the GBNF grammar string and json_mode flag from a request.

    Returns (grammar_string_or_None, is_json_mode).
    """
    from .grammar import json_mode_grammar, json_schema_to_gbnf

    # Explicit grammar takes precedence
    if body.grammar:
        return body.grammar, False

    rf = body.response_format
    if rf is None and default_json_mode:
        rf = {"type": "json_object"}

    if rf is None:
        return None, False

    fmt_type = rf.get("type")
    if fmt_type == "json_object":
        return json_mode_grammar(), True
    if fmt_type == "json_schema":
        schema = rf.get("json_schema") or rf.get("schema")
        if schema:
            # The OpenAI API wraps the actual schema under a "schema" key
            # inside json_schema. Handle both nesting levels.
            actual_schema = schema.get("schema", schema)
            return json_schema_to_gbnf(actual_schema), True
        return json_mode_grammar(), True

    return None, False


def _inject_json_system_prompt(
    messages: list[dict[str, str]],
    schema: Optional[dict[str, Any]] = None,
) -> list[dict[str, str]]:
    """Prepend a JSON-mode system prompt if one isn't already present."""
    from .grammar import json_system_prompt

    # Don't double-inject
    if messages and messages[0].get("role") == "system":
        existing = messages[0].get("content", "")
        if "JSON" in existing or "json" in existing:
            return messages

    prompt = json_system_prompt(schema)
    return [{"role": "system", "content": prompt}] + list(messages)


def create_app(
    model_name: str,
    *,
    api_key: Optional[str] = None,
    api_base: str = "https://api.octomil.com/api/v1",
    json_mode: bool = False,
    cache_size_mb: int = 2048,
    cache_enabled: bool = True,
    engine: Optional[str] = None,
    max_queue_depth: int = 32,
    moe_config: Optional[MoEConfig] = None,
    compress_context: bool = False,
    compression_strategy: str = "token_pruning",
    compression_ratio: float = 0.5,
    compression_max_turns: int = 4,
    compression_threshold: int = 256,
    tool_use: bool = False,
    early_exit_config: Optional["EarlyExitConfig"] = None,
) -> Any:
    """Create a FastAPI app with OpenAI-compatible endpoints.

    Parameters
    ----------
    json_mode:
        When ``True``, all requests default to ``response_format={"type": "json_object"}``
        unless the caller explicitly sets a different ``response_format``.
    engine:
        Force a specific engine (e.g. ``"mlx-lm"``, ``"llama.cpp"``).
        When ``None``, auto-benchmarks all available engines and picks fastest.
    max_queue_depth:
        Maximum number of pending requests in the queue.  When full, new
        requests get a 503 response.  Set to 0 to disable queueing.
    moe_config:
        Configuration for Mixture of Experts model features.
        When ``None``, uses defaults (auto-detection enabled).
    compress_context:
        Enable prompt compression.  Long prompts are compressed before
        inference to reduce context window usage and speed up prefill.
    compression_strategy:
        Compression strategy: ``"token_pruning"`` or ``"sliding_window"``.
    compression_ratio:
        Target compression ratio (0.0--1.0) for token pruning.
    compression_max_turns:
        Number of recent turns to keep verbatim (sliding_window strategy).
    compression_threshold:
        Minimum estimated token count before compression activates.
    tool_use:
        When ``True``, pre-load coding agent tool schemas and expose
        them at ``/v1/tool-schemas``.  Tools include ``read_file``,
        ``write_file``, ``edit_file``, ``run_command``, ``search_files``.
    early_exit_config:
        Configuration for early exit / adaptive computation depth.
        When ``None`` or not enabled, early exit monitoring is disabled.
    """
    from contextlib import asynccontextmanager

    from fastapi import FastAPI, HTTPException
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import StreamingResponse

    # Detect MoE model from catalog
    from .models.catalog import get_moe_metadata
    from .models.catalog import is_moe_model as _is_moe

    _moe_detected = _is_moe(model_name)
    _moe_meta = get_moe_metadata(model_name)

    # Build compressor if enabled
    _compressor = None
    if compress_context:
        from .compression import CompressionConfig, PromptCompressor

        _compressor = PromptCompressor(
            CompressionConfig(
                enabled=True,
                strategy=compression_strategy,
                target_ratio=compression_ratio,
                max_turns_verbatim=compression_max_turns,
                token_threshold=compression_threshold,
            )
        )

    state = ServerState(
        model_name=model_name,
        api_key=api_key,
        api_base=api_base,
        default_json_mode=json_mode,
        cache_size_mb=cache_size_mb,
        cache_enabled=cache_enabled,
        engine_override=engine,
        max_queue_depth=max_queue_depth,
        moe_config=moe_config or MoEConfig(),
        is_moe_model=_moe_detected,
        moe_metadata=_moe_meta,
        compressor=_compressor,
        early_exit_config=early_exit_config,
        tool_use=tool_use,
    )

    @asynccontextmanager
    async def lifespan(app: Any) -> Any:
        if state.is_moe_model and state.moe_metadata and state.moe_config.enabled:
            _m = state.moe_metadata
            logger.info(
                "MoE model detected: %s (%d experts, %d active per token, " "%s total params, %s active params)",
                model_name,
                _m.num_experts,
                _m.active_experts,
                _m.total_params,
                _m.active_params,
            )

        # Check if this is a whisper (speech-to-text) model
        from .engines.whisper_engine import is_whisper_model

        if is_whisper_model(model_name):
            from .engines.whisper_engine import WhisperCppEngine

            whisper_engine = WhisperCppEngine()
            whisper_backend = whisper_engine.create_backend(model_name)
            try:
                whisper_backend.load_model(model_name)
            except Exception as exc:
                _log_startup_error(model_name, exc)
                raise
            state.whisper_backend = whisper_backend
            state.engine_name = "whisper.cpp"
        else:
            try:
                state.backend = _detect_backend(
                    model_name,
                    cache_size_mb=state.cache_size_mb,
                    cache_enabled=state.cache_enabled,
                    engine_override=state.engine_override,
                )
            except Exception as exc:
                _log_startup_error(model_name, exc)
                raise
            state.engine_name = state.backend.name if state.backend else "none"
        state.start_time = time.time()

        # Initialise early exit monitor
        if state.early_exit_config is not None and state.early_exit_config.enabled:
            from .early_exit import EarlyExitMonitor as _EEM

            state.early_exit_monitor = _EEM(config=state.early_exit_config)
            logger.info(
                "Early exit enabled (threshold=%.2f, min_layers_frac=%.2f%s)",
                state.early_exit_config.effective_threshold,
                state.early_exit_config.effective_min_layers_fraction,
                f", preset={state.early_exit_config.preset.value}" if state.early_exit_config.preset else "",
            )

        # Initialise request queue
        if state.max_queue_depth > 0:
            from .batch import RequestQueue

            state.request_queue = RequestQueue(max_depth=state.max_queue_depth)
            state.request_queue.start()
            logger.info("Request queue enabled (max_depth=%d)", state.max_queue_depth)

        # Create telemetry reporter when an API key is configured
        if state.api_key:
            try:
                from .telemetry import TelemetryReporter as _TR

                state.reporter = _TR(
                    api_key=state.api_key,
                    api_base=state.api_base,
                    org_id="default",
                )
                logger.info("Telemetry enabled — reporting to %s", state.api_base)
            except Exception as exc:
                logger.warning("Failed to initialise telemetry: %s", exc)

        yield

        # Graceful shutdown: stop request queue
        if state.request_queue is not None:
            await state.request_queue.stop()

        # Graceful shutdown: drain pending telemetry events
        if state.reporter is not None:
            state.reporter.close()

    app = FastAPI(title="Octomil Serve", version="1.0.0", lifespan=lifespan)

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @app.middleware("http")
    async def agent_context_middleware(request, call_next):  # type: ignore[no-untyped-def]
        """Log and store agent context from the X-Octomil-Agent-Context header.

        Coding agents (Aider, Goose, OpenCode) can send this header to
        identify themselves and provide context for smarter routing decisions.
        """
        agent_ctx = request.headers.get("X-Octomil-Agent-Context")
        if agent_ctx:
            logger.info("Agent context: %s", agent_ctx)
            # Store on request state for downstream use
            request.state.agent_context = agent_ctx
        else:
            request.state.agent_context = None
        response = await call_next(request)
        return response

    @app.get("/v1/models")
    async def list_models() -> dict[str, Any]:
        models = state.backend.list_models() if state.backend else []
        data = []
        for m in models:
            entry: dict[str, Any] = {
                "id": m,
                "object": "model",
                "created": int(state.start_time),
                "owned_by": "octomil",
            }
            if state.is_moe_model and state.moe_metadata:
                entry["architecture"] = "moe"
                entry["moe"] = {
                    "num_experts": state.moe_metadata.num_experts,
                    "active_experts": state.moe_metadata.active_experts,
                    "expert_size": state.moe_metadata.expert_size,
                    "total_params": state.moe_metadata.total_params,
                    "active_params": state.moe_metadata.active_params,
                }
            else:
                entry["architecture"] = "dense"
            data.append(entry)
        return {
            "object": "list",
            "data": data,
        }

    @app.get("/v1/tool-schemas")
    async def tool_schemas() -> dict[str, Any]:
        """Return pre-loaded coding agent tool schemas.

        Only populated when the server is started with ``--tool-use``.
        Coding agents (Aider, Goose, OpenCode) can discover tools here
        and use them for structured output enforcement.
        """
        if not state.tool_use:
            return {"tools": [], "enabled": False}

        from .tool_schemas import get_tool_use_tools

        return {"tools": get_tool_use_tools(), "enabled": True}

    @app.post("/v1/chat/completions")
    async def chat_completions(body: ChatCompletionBody) -> Any:
        if state.backend is None:
            raise HTTPException(status_code=503, detail="Model not loaded")

        # --- Tier-0: deterministic routing (arithmetic, etc.) ---
        # Check the last user message for a query that can be answered
        # without invoking any model.
        if body.messages and not body.stream:
            last_user_msg = ""
            for msg in reversed(body.messages):
                if msg.role == "user":
                    last_user_msg = msg.content
                    break
            if last_user_msg:
                from .routing import check_deterministic

                det = check_deterministic(last_user_msg)
                if det is not None:
                    state.request_count += 1
                    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
                    return {
                        "id": req_id,
                        "object": "chat.completion",
                        "created": int(time.time()),
                        "model": body.model or state.model_name,
                        "choices": [
                            {
                                "index": 0,
                                "message": {
                                    "role": "assistant",
                                    "content": det.answer,
                                },
                                "finish_reason": "stop",
                            }
                        ],
                        "usage": {
                            "prompt_tokens": 0,
                            "completion_tokens": 0,
                            "total_tokens": 0,
                            "deterministic": True,
                            "deterministic_method": det.method,
                        },
                    }

        grammar_str, is_json = _resolve_grammar(body, state.default_json_mode)

        messages = [{"role": m.role, "content": m.content} for m in body.messages]

        # --- Prompt compression ---
        compression_stats = None
        if state.compressor is not None:
            messages, compression_stats = state.compressor.compress(messages)

        # For backends without native grammar support (MLX, echo),
        # inject a system prompt nudging JSON output when json_mode is on.
        uses_grammar_natively = isinstance(state.backend, LlamaCppBackend)
        schema_for_prompt: Optional[dict[str, Any]] = None
        if is_json and not uses_grammar_natively:
            rf = body.response_format or {}
            if rf.get("type") == "json_schema":
                raw = rf.get("json_schema") or rf.get("schema")
                schema_for_prompt = raw.get("schema", raw) if raw else None
            messages = _inject_json_system_prompt(messages, schema_for_prompt)

        gen_req = GenerationRequest(
            model=body.model or state.model_name,
            messages=messages,
            max_tokens=body.max_tokens,
            temperature=body.temperature,
            top_p=body.top_p,
            stream=body.stream,
            grammar=grammar_str if uses_grammar_natively else None,
            json_mode=is_json,
        )

        state.request_count += 1
        req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        session_id = uuid.uuid4().hex
        model_version = "latest"
        _reporter = state.reporter

        # Report inference_started (best-effort)
        if _reporter is not None:
            try:
                _reporter.report_inference_started(
                    model_id=gen_req.model,
                    version=model_version,
                    session_id=session_id,
                )
            except Exception:
                pass

        # Report compression telemetry (best-effort)
        if _reporter is not None and compression_stats is not None:
            try:
                _reporter.report_prompt_compressed(
                    session_id=session_id,
                    model_id=gen_req.model,
                    version=model_version,
                    original_tokens=compression_stats.original_tokens,
                    compressed_tokens=compression_stats.compressed_tokens,
                    compression_ratio=compression_stats.compression_ratio,
                    strategy=compression_stats.strategy,
                    duration_ms=compression_stats.duration_ms,
                )
            except Exception:
                pass

        # --- Queue-aware dispatch ---
        _queue = state.request_queue

        if gen_req.stream:
            if _queue is not None:
                # Stream through the request queue
                return StreamingResponse(
                    _queued_stream_response(state, gen_req, req_id, session_id, _queue),
                    media_type="text/event-stream",
                )
            return StreamingResponse(
                _stream_response(state, gen_req, req_id, session_id),
                media_type="text/event-stream",
            )

        gen_start = time.monotonic()
        try:
            if _queue is not None:
                from .batch import QueueFullError, QueueTimeoutError

                try:
                    text, metrics = await _queue.submit_generate(gen_req, state.backend.generate)
                except QueueFullError:
                    raise HTTPException(
                        status_code=503,
                        detail="Server busy — request queue full. Try again later.",
                    )
                except QueueTimeoutError:
                    raise HTTPException(
                        status_code=504,
                        detail="Request timed out waiting in queue.",
                    )
            else:
                text, metrics = state.backend.generate(gen_req)
        except HTTPException:
            raise
        except Exception:
            if _reporter is not None:
                try:
                    _reporter.report_inference_failed(
                        session_id=session_id,
                        model_id=gen_req.model,
                        version=model_version,
                    )
                except Exception:
                    pass
            raise
        gen_elapsed_ms = (time.monotonic() - gen_start) * 1000

        # JSON validation + retry for non-grammar backends
        if is_json and not uses_grammar_natively:
            from .grammar import extract_json, validate_json_output

            if not validate_json_output(text):
                extracted = extract_json(text)
                if extracted is not None:
                    text = json.dumps(extracted)
                else:
                    # Retry once with a stronger system prompt
                    retry_messages = _inject_json_system_prompt(messages, schema_for_prompt)
                    retry_req = GenerationRequest(
                        model=gen_req.model,
                        messages=retry_messages,
                        max_tokens=gen_req.max_tokens,
                        temperature=max(gen_req.temperature - 0.2, 0.0),
                        top_p=gen_req.top_p,
                        stream=False,
                        json_mode=True,
                    )
                    text, metrics = state.backend.generate(retry_req)
                    # Best-effort extraction on retry
                    if not validate_json_output(text):
                        extracted = extract_json(text)
                        if extracted is not None:
                            text = json.dumps(extracted)

        # Early exit monitoring (best-effort)
        _ee_monitor = state.early_exit_monitor
        ee_metrics_dict: Optional[dict[str, Any]] = None
        if _ee_monitor is not None and _ee_monitor.config.enabled:
            try:
                total_layers = _ee_monitor.config.total_layers or 32
                ee_req_metrics = _ee_monitor.simulate_token_exits(
                    token_count=metrics.total_tokens,
                    total_layers=total_layers,
                )
                _ee_monitor.record_request(ee_req_metrics)
                metrics.early_exit_tokens = ee_req_metrics.early_exit_tokens
                metrics.avg_layers_used = ee_req_metrics.avg_layers_used
                ee_metrics_dict = ee_req_metrics.to_dict()
            except Exception:
                pass

        # Report inference_completed (best-effort)
        if _reporter is not None:
            try:
                total_tokens = metrics.total_tokens
                throughput = total_tokens / (gen_elapsed_ms / 1000) if gen_elapsed_ms > 0 else 0.0
                _reporter.report_inference_completed(
                    session_id=session_id,
                    model_id=gen_req.model,
                    version=model_version,
                    total_chunks=total_tokens,
                    total_duration_ms=gen_elapsed_ms,
                    ttfc_ms=metrics.ttfc_ms,
                    throughput=throughput,
                    attention_backend=metrics.attention_backend,
                    early_exit_stats=ee_metrics_dict,
                )
            except Exception:
                pass

        usage: dict[str, Any] = {
            "prompt_tokens": metrics.prompt_tokens,
            "completion_tokens": metrics.total_tokens,
            "total_tokens": metrics.prompt_tokens + metrics.total_tokens,
            "cache_hit": metrics.cache_hit,
        }

        # Include compression stats when compression was applied
        if compression_stats is not None and compression_stats.strategy != "none":
            usage["compression"] = {
                "original_tokens": compression_stats.original_tokens,
                "compressed_tokens": compression_stats.compressed_tokens,
                "ratio": round(compression_stats.compression_ratio, 4),
                "strategy": compression_stats.strategy,
                "duration_ms": round(compression_stats.duration_ms, 2),
            }
        if ee_metrics_dict is not None:
            usage["early_exit"] = ee_metrics_dict

        return {
            "id": req_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": gen_req.model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }
            ],
            "usage": usage,
        }

    @app.get("/v1/cache/stats")
    async def cache_stats() -> dict[str, Any]:
        """Return KV cache statistics."""
        if state.backend is None:
            raise HTTPException(status_code=503, detail="Model not loaded")

        cache_mgr = _get_cache_manager(state.backend)
        if cache_mgr is None:
            return {
                "enabled": False,
                "hits": 0,
                "misses": 0,
                "hit_rate": 0.0,
                "entries": 0,
                "memory_mb": 0.0,
                "max_memory_mb": state.cache_size_mb,
            }

        stats = cache_mgr.stats()
        return {
            "enabled": True,
            "hits": stats.hits,
            "misses": stats.misses,
            "hit_rate": round(stats.hit_rate, 4),
            "entries": stats.entries,
            "memory_mb": round(stats.memory_mb, 2),
            "max_memory_mb": state.cache_size_mb,
        }

    @app.get("/v1/queue/stats")
    async def queue_stats() -> dict[str, Any]:
        """Return request queue statistics."""
        if state.request_queue is None:
            return {
                "pending": 0,
                "active": 0,
                "max_depth": 0,
                "enabled": False,
            }
        qs = state.request_queue.stats()
        return {
            "pending": qs.pending,
            "active": qs.active,
            "max_depth": qs.max_depth,
            "enabled": True,
        }

    @app.get("/v1/early-exit/stats")
    async def early_exit_stats() -> dict[str, Any]:
        """Return early exit / adaptive computation depth statistics."""
        if state.early_exit_monitor is None:
            return {
                "enabled": False,
                "config": {},
                "stats": {
                    "total_requests": 0,
                    "total_tokens": 0,
                    "total_early_exit_tokens": 0,
                    "exit_percentage": 0.0,
                    "avg_layers_used": 0.0,
                    "avg_entropy": 0.0,
                },
            }
        return state.early_exit_monitor.get_stats_dict()

    @app.get("/v1/engines")
    async def list_engines() -> dict[str, Any]:
        """List detected engines and their benchmark results."""
        from .engines import get_registry

        registry = get_registry()
        detections = registry.detect_all(state.model_name)

        engines_list = []
        for d in detections:
            entry: dict[str, Any] = {
                "name": d.engine.name,
                "display_name": d.engine.display_name,
                "available": d.available,
                "info": d.info,
                "active": (state.backend is not None and state.backend.name == d.engine.name),
            }
            engines_list.append(entry)

        return {
            "active_engine": state.engine_name,
            "engines": engines_list,
        }

    @app.get("/health")
    async def health() -> dict[str, Any]:
        cache_info: dict[str, Any] = {"enabled": state.cache_enabled}
        cache_mgr = _get_cache_manager(state.backend) if state.backend is not None else None
        if cache_mgr is not None:
            cs = cache_mgr.stats()
            cache_info["entries"] = cs.entries
            cache_info["hit_rate"] = round(cs.hit_rate, 4)

        backend_name = "none"
        if state.whisper_backend is not None:
            backend_name = state.whisper_backend.name
        elif state.backend is not None:
            backend_name = state.backend.name

        health_data: dict[str, Any] = {
            "status": "ok",
            "model": state.model_name,
            "engine": state.engine_name,
            "backend": backend_name,
            "requests_served": state.request_count,
            "uptime_seconds": int(time.time() - state.start_time),
            "cache": cache_info,
        }
        if state.early_exit_monitor is not None:
            ee_stats = state.early_exit_monitor.stats
            health_data["early_exit"] = {
                "enabled": True,
                "exit_percentage": round(ee_stats.exit_percentage, 2),
                "avg_layers_used": round(ee_stats.avg_layers_used, 2),
            }
        return health_data

    @app.post("/v1/audio/transcriptions")
    async def transcribe_audio(
        file: Any = None,
        model: str = "",
    ) -> dict[str, Any]:
        """OpenAI Whisper API-compatible audio transcription endpoint.

        Accepts a multipart file upload and returns transcribed text
        with segment-level timestamps.
        """
        if state.whisper_backend is None:
            raise HTTPException(
                status_code=503,
                detail="No whisper model loaded. Start server with a whisper model: " "octomil serve whisper-base",
            )

        if file is None:
            raise HTTPException(status_code=400, detail="No audio file provided")

        # Save uploaded file to a temp location
        import tempfile

        filename = getattr(file, "filename", "audio.wav") or "audio.wav"
        suffix = os.path.splitext(filename)[1] or ".wav"
        fd, tmp_path = tempfile.mkstemp(suffix=suffix)
        try:
            content = await file.read()
            os.write(fd, content)
            os.close(fd)

            state.request_count += 1
            result: dict[str, Any] = state.whisper_backend.transcribe(tmp_path)
            return result
        finally:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    return app


async def _stream_response(
    state: Union[ServerState, "_MultiModelStateAdapter"],
    request: GenerationRequest,
    req_id: str,
    session_id: str = "",
) -> AsyncIterator[str]:
    """Yield SSE chunks in OpenAI streaming format."""
    assert state.backend is not None

    _reporter = state.reporter
    model_version = "latest"
    chunk_index = 0
    stream_start = time.monotonic()
    first_chunk_time: Optional[float] = None
    prev_chunk_time = stream_start
    failed = False

    try:
        async for chunk in state.backend.generate_stream(request):
            now = time.monotonic()
            chunk_latency_ms = (now - prev_chunk_time) * 1000
            prev_chunk_time = now

            if first_chunk_time is None and chunk.text:
                first_chunk_time = now

            # Report each chunk (best-effort)
            if _reporter is not None and chunk.text:
                try:
                    ttfc = (
                        (first_chunk_time - stream_start) * 1000
                        if first_chunk_time is not None and chunk_index == 0
                        else None
                    )
                    _reporter.report_inference_chunk(
                        session_id=session_id,
                        model_id=request.model,
                        version=model_version,
                        chunk_index=chunk_index,
                        ttfc_ms=ttfc,
                        chunk_latency_ms=chunk_latency_ms,
                    )
                except Exception:
                    pass
                chunk_index += 1

            data = {
                "id": req_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": request.model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": chunk.text} if chunk.text else {},
                        "finish_reason": chunk.finish_reason,
                    }
                ],
            }
            yield f"data: {json.dumps(data)}\n\n"
    except Exception:
        failed = True
        if _reporter is not None:
            try:
                _reporter.report_inference_failed(
                    session_id=session_id,
                    model_id=request.model,
                    version=model_version,
                )
            except Exception:
                pass
        raise

    yield "data: [DONE]\n\n"

    # Report inference_completed (best-effort)
    if _reporter is not None and not failed:
        try:
            total_duration_ms = (time.monotonic() - stream_start) * 1000
            ttfc_ms = (first_chunk_time - stream_start) * 1000 if first_chunk_time is not None else 0.0
            throughput = chunk_index / (total_duration_ms / 1000) if total_duration_ms > 0 else 0.0
            _reporter.report_inference_completed(
                session_id=session_id,
                model_id=request.model,
                version=model_version,
                total_chunks=chunk_index,
                total_duration_ms=total_duration_ms,
                ttfc_ms=ttfc_ms,
                throughput=throughput,
                attention_backend=(state.backend.attention_backend if state.backend is not None else None),
            )
        except Exception:
            pass


async def _queued_stream_response(
    state: Any,
    request: GenerationRequest,
    req_id: str,
    session_id: str,
    queue: Any,
) -> AsyncIterator[str]:
    """Yield SSE chunks after waiting in the request queue.

    Submits a streaming request to the queue.  Once the request reaches
    the front, chunks are forwarded as SSE events.  Queue errors (full,
    timeout) are surfaced as SSE error events so the client gets useful
    feedback even on a streaming connection.
    """
    from .batch import QueueFullError, QueueTimeoutError

    assert state.backend is not None

    _reporter = state.reporter
    model_version = "latest"
    stream_start = time.monotonic()
    first_chunk_time: Optional[float] = None
    prev_chunk_time = stream_start

    try:
        chunk_iter = queue.submit_generate_stream(request, state.backend.generate_stream)
        # Build SSE events from the chunk iterator (same format as _stream_response)
        chunk_index = 0
        async for chunk in chunk_iter:
            now = time.monotonic()
            chunk_latency_ms = (now - prev_chunk_time) * 1000
            prev_chunk_time = now

            if first_chunk_time is None and chunk.text:
                first_chunk_time = now

            # Report each chunk (best-effort)
            if _reporter is not None and chunk.text:
                try:
                    ttfc = (
                        (first_chunk_time - stream_start) * 1000
                        if first_chunk_time is not None and chunk_index == 0
                        else None
                    )
                    _reporter.report_inference_chunk(
                        session_id=session_id,
                        model_id=request.model,
                        version=model_version,
                        chunk_index=chunk_index,
                        ttfc_ms=ttfc,
                        chunk_latency_ms=chunk_latency_ms,
                    )
                except Exception:
                    pass

            data = {
                "id": req_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": request.model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": chunk.text} if chunk.text else {},
                        "finish_reason": chunk.finish_reason,
                    }
                ],
            }
            yield f"data: {json.dumps(data)}\n\n"
            chunk_index += 1

        yield "data: [DONE]\n\n"

        # Report inference_completed (best-effort)
        if _reporter is not None:
            try:
                total_duration_ms = (time.monotonic() - stream_start) * 1000
                ttfc_ms = (first_chunk_time - stream_start) * 1000 if first_chunk_time is not None else 0.0
                throughput = chunk_index / (total_duration_ms / 1000) if total_duration_ms > 0 else 0.0
                _reporter.report_inference_completed(
                    session_id=session_id,
                    model_id=request.model,
                    version=model_version,
                    total_chunks=chunk_index,
                    total_duration_ms=total_duration_ms,
                    ttfc_ms=ttfc_ms,
                    throughput=throughput,
                )
            except Exception:
                pass
    except QueueFullError:
        error_data = {
            "error": {
                "message": "Server busy — request queue full. Try again later.",
                "type": "server_error",
                "code": 503,
            }
        }
        yield f"data: {json.dumps(error_data)}\n\n"
    except QueueTimeoutError:
        error_data = {
            "error": {
                "message": "Request timed out waiting in queue.",
                "type": "server_error",
                "code": 504,
            }
        }
        yield f"data: {json.dumps(error_data)}\n\n"


def run_server(
    model_name: str,
    *,
    port: int = 8080,
    host: str = "0.0.0.0",
    api_key: Optional[str] = None,
    api_base: str = "https://api.octomil.com/api/v1",
    json_mode: bool = False,
    cache_size_mb: int = 2048,
    cache_enabled: bool = True,
    engine: Optional[str] = None,
    max_queue_depth: int = 32,
    moe_config: Optional[MoEConfig] = None,
    compress_context: bool = False,
    compression_strategy: str = "token_pruning",
    compression_ratio: float = 0.5,
    compression_max_turns: int = 4,
    compression_threshold: int = 256,
    tool_use: bool = False,
    early_exit_config: Optional["EarlyExitConfig"] = None,
) -> None:
    """Start the inference server (blocking).

    Parameters
    ----------
    json_mode:
        When ``True``, all requests default to ``response_format={"type": "json_object"}``.
    engine:
        Force a specific engine (e.g. ``"mlx-lm"``, ``"llama.cpp"``).
        When ``None``, auto-benchmarks and picks fastest.
    max_queue_depth:
        Maximum number of pending requests in the queue (default 32).
        Set to 0 to disable queueing.
    moe_config:
        Configuration for Mixture of Experts model features.
        When ``None``, uses defaults (auto-detection enabled).
    compress_context:
        Enable prompt compression before inference.
    compression_strategy:
        Compression strategy: ``"token_pruning"`` or ``"sliding_window"``.
    compression_ratio:
        Target compression ratio (0.0--1.0) for token pruning.
    compression_max_turns:
        Number of recent turns to keep verbatim (sliding_window).
    compression_threshold:
        Minimum token count before compression activates.
    tool_use:
        Pre-load coding agent tool schemas for structured output.
    early_exit_config:
        Configuration for early exit / adaptive computation depth.
    """
    import uvicorn

    app = create_app(
        model_name,
        api_key=api_key,
        api_base=api_base,
        json_mode=json_mode,
        cache_size_mb=cache_size_mb,
        cache_enabled=cache_enabled,
        engine=engine,
        max_queue_depth=max_queue_depth,
        moe_config=moe_config,
        compress_context=compress_context,
        compression_strategy=compression_strategy,
        compression_ratio=compression_ratio,
        compression_max_turns=compression_max_turns,
        compression_threshold=compression_threshold,
        tool_use=tool_use,
        early_exit_config=early_exit_config,
    )
    uvicorn.run(app, host=host, port=port, log_level="info")


# ---------------------------------------------------------------------------
# Multi-model serving with query routing
# ---------------------------------------------------------------------------


def create_multi_model_app(
    model_names: list[str],
    *,
    api_key: Optional[str] = None,
    api_base: str = "https://api.octomil.com/api/v1",
    json_mode: bool = False,
    cache_size_mb: int = 2048,
    cache_enabled: bool = True,
    engine: Optional[str] = None,
    route_strategy: str = "complexity",
) -> Any:
    """Create a FastAPI app that loads multiple models and routes queries.

    Parameters
    ----------
    model_names:
        Ordered list of model names, smallest to largest.
    route_strategy:
        Routing strategy (currently only ``"complexity"``).
    """
    from contextlib import asynccontextmanager

    from fastapi import FastAPI, HTTPException
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import JSONResponse, StreamingResponse

    from .decomposer import ResultMerger, SubTaskResult
    from .routing import (
        DecomposedRoutingDecision,
        QueryRouter,
        RoutingDecision,
        assign_tiers,
    )

    state = MultiModelServerState(
        model_names=model_names,
        api_key=api_key,
        api_base=api_base,
        default_json_mode=json_mode,
        cache_size_mb=cache_size_mb,
        cache_enabled=cache_enabled,
        engine_override=engine,
        route_strategy=route_strategy,
    )

    @asynccontextmanager
    async def lifespan(app: Any) -> Any:
        # Load all models
        for name in model_names:
            try:
                backend = _detect_backend(
                    name,
                    cache_size_mb=state.cache_size_mb,
                    cache_enabled=state.cache_enabled,
                    engine_override=state.engine_override,
                )
                state.backends[name] = backend
                state.routed_counts[name] = 0
                logger.info("Loaded model: %s (engine: %s)", name, backend.name)
            except Exception as exc:
                _log_startup_error(name, exc)
                raise

        # Build router with auto-assigned tiers
        model_infos = assign_tiers(model_names)
        state.router = QueryRouter(
            model_infos,
            strategy=state.route_strategy,
        )

        state.start_time = time.time()

        # Create telemetry reporter
        if state.api_key:
            try:
                from .telemetry import TelemetryReporter as _TR

                state.reporter = _TR(
                    api_key=state.api_key,
                    api_base=state.api_base,
                    org_id="default",
                )
            except Exception as exc:
                logger.warning("Failed to initialise telemetry: %s", exc)

        yield

        if state.reporter is not None:
            state.reporter.close()

    app = FastAPI(title="Octomil Serve (Multi-Model)", version="1.0.0", lifespan=lifespan)

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @app.get("/v1/models")
    async def list_models() -> dict[str, Any]:
        all_models: list[str] = []
        for backend in state.backends.values():
            all_models.extend(backend.list_models())
        return {
            "object": "list",
            "data": [
                {
                    "id": m,
                    "object": "model",
                    "created": int(state.start_time),
                    "owned_by": "octomil",
                }
                for m in all_models
            ],
        }

    @app.post("/v1/chat/completions")
    async def chat_completions(body: ChatCompletionBody) -> Any:
        if not state.backends:
            raise HTTPException(status_code=503, detail="No models loaded")

        messages = [{"role": m.role, "content": m.content} for m in body.messages]

        # Attempt query decomposition for non-streaming requests
        assert state.router is not None
        if not body.stream:
            decomp_decision = state.router.route_decomposed(messages)
            if isinstance(decomp_decision, DecomposedRoutingDecision):
                return await _handle_decomposed(state, body, messages, decomp_decision)

        # Standard single-task routing
        decision: RoutingDecision = state.router.route(messages)

        routed_model = decision.model_name
        fallback_chain = decision.fallback_chain

        # Try the routed model, then fallback chain
        last_error: Optional[Exception] = None
        tried_models: list[str] = [routed_model] + fallback_chain
        used_fallback = False

        for model_name in tried_models:
            backend = state.backends.get(model_name)
            if backend is None:
                continue

            grammar_str, is_json = _resolve_grammar(body, state.default_json_mode)

            req_messages = list(messages)
            uses_grammar_natively = isinstance(backend, LlamaCppBackend)
            schema_for_prompt: Optional[dict[str, Any]] = None
            if is_json and not uses_grammar_natively:
                rf = body.response_format or {}
                if rf.get("type") == "json_schema":
                    raw = rf.get("json_schema") or rf.get("schema")
                    schema_for_prompt = raw.get("schema", raw) if raw else None
                req_messages = _inject_json_system_prompt(req_messages, schema_for_prompt)

            gen_req = GenerationRequest(
                model=model_name,
                messages=req_messages,
                max_tokens=body.max_tokens,
                temperature=body.temperature,
                top_p=body.top_p,
                stream=body.stream,
                grammar=grammar_str if uses_grammar_natively else None,
                json_mode=is_json,
            )

            state.request_count += 1
            state.routed_counts[model_name] = state.routed_counts.get(model_name, 0) + 1
            req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
            session_id = uuid.uuid4().hex
            model_version = "latest"
            _reporter = state.reporter

            # Report inference_started
            if _reporter is not None:
                try:
                    _reporter.report_inference_started(
                        model_id=model_name,
                        version=model_version,
                        session_id=session_id,
                    )
                except Exception:
                    pass

            if gen_req.stream:
                headers = {
                    "X-Octomil-Routed-Model": model_name,
                    "X-Octomil-Complexity": str(decision.complexity_score),
                    "X-Octomil-Tier": decision.tier,
                }
                if used_fallback:
                    headers["X-Octomil-Fallback"] = "true"

                return StreamingResponse(
                    _stream_response(
                        _MultiModelStateAdapter(state, backend, model_name),
                        gen_req,
                        req_id,
                        session_id,
                    ),
                    media_type="text/event-stream",
                    headers=headers,
                )

            gen_start = time.monotonic()
            try:
                text, metrics = backend.generate(gen_req)
            except Exception as exc:
                last_error = exc
                used_fallback = True
                state.fallback_counts += 1
                logger.warning("Model %s failed, trying fallback: %s", model_name, exc)
                if _reporter is not None:
                    try:
                        _reporter.report_inference_failed(
                            session_id=session_id,
                            model_id=model_name,
                            version=model_version,
                        )
                    except Exception:
                        pass
                continue
            gen_elapsed_ms = (time.monotonic() - gen_start) * 1000

            # JSON validation + retry for non-grammar backends
            if is_json and not uses_grammar_natively:
                from .grammar import extract_json, validate_json_output

                if not validate_json_output(text):
                    extracted = extract_json(text)
                    if extracted is not None:
                        text = json.dumps(extracted)
                    else:
                        retry_messages = _inject_json_system_prompt(messages, schema_for_prompt)
                        retry_req = GenerationRequest(
                            model=model_name,
                            messages=retry_messages,
                            max_tokens=gen_req.max_tokens,
                            temperature=max(gen_req.temperature - 0.2, 0.0),
                            top_p=gen_req.top_p,
                            stream=False,
                            json_mode=True,
                        )
                        text, metrics = backend.generate(retry_req)
                        if not validate_json_output(text):
                            extracted = extract_json(text)
                            if extracted is not None:
                                text = json.dumps(extracted)

            # Report inference_completed
            if _reporter is not None:
                try:
                    total_tokens = metrics.total_tokens
                    throughput = total_tokens / (gen_elapsed_ms / 1000) if gen_elapsed_ms > 0 else 0.0
                    _reporter.report_inference_completed(
                        session_id=session_id,
                        model_id=model_name,
                        version=model_version,
                        total_chunks=total_tokens,
                        total_duration_ms=gen_elapsed_ms,
                        ttfc_ms=metrics.ttfc_ms,
                        throughput=throughput,
                        attention_backend=metrics.attention_backend,
                    )
                except Exception:
                    pass

            response_data = {
                "id": req_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model_name,
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": text},
                        "finish_reason": "stop",
                    }
                ],
                "usage": {
                    "prompt_tokens": metrics.prompt_tokens,
                    "completion_tokens": metrics.total_tokens,
                    "total_tokens": metrics.prompt_tokens + metrics.total_tokens,
                    "cache_hit": metrics.cache_hit,
                },
            }

            resp = JSONResponse(content=response_data)
            resp.headers["X-Octomil-Routed-Model"] = model_name
            resp.headers["X-Octomil-Complexity"] = str(decision.complexity_score)
            resp.headers["X-Octomil-Tier"] = decision.tier
            if used_fallback:
                resp.headers["X-Octomil-Fallback"] = "true"
            return resp

        # All models failed
        raise HTTPException(
            status_code=503,
            detail=f"All models failed. Last error: {last_error}",
        )

    @app.get("/v1/routing/stats")
    async def routing_stats() -> dict[str, Any]:
        """Return routing statistics."""
        return {
            "total_requests": state.request_count,
            "routed_counts": dict(state.routed_counts),
            "fallback_count": state.fallback_counts,
            "models": state.model_names,
            "strategy": state.route_strategy,
        }

    @app.get("/health")
    async def health() -> dict[str, Any]:
        models_status = {}
        for name, backend in state.backends.items():
            models_status[name] = {
                "engine": backend.name,
                "requests": state.routed_counts.get(name, 0),
            }
        return {
            "status": "ok",
            "mode": "multi-model",
            "models": models_status,
            "strategy": state.route_strategy,
            "requests_served": state.request_count,
            "fallback_count": state.fallback_counts,
            "uptime_seconds": int(time.time() - state.start_time),
        }

    async def _handle_decomposed(
        mm_state: MultiModelServerState,
        body: ChatCompletionBody,
        messages: list[dict[str, str]],
        decomp: DecomposedRoutingDecision,
    ) -> Any:
        """Execute a decomposed multi-task query with parallel dispatch.

        Independent sub-tasks are executed concurrently via asyncio.gather.
        Dependent sub-tasks wait for their prerequisites before executing.
        """
        import asyncio

        decomposer_merger = ResultMerger()
        sub_task_results: dict[int, SubTaskResult] = {}

        async def _execute_subtask(
            task_idx: int,
        ) -> SubTaskResult:
            task = decomp.tasks[task_idx]
            decision = decomp.sub_decisions[task_idx]
            model_name = decision.model_name
            backend = mm_state.backends.get(model_name)

            if backend is None:
                # Use first available backend as fallback
                model_name = next(iter(mm_state.backends))
                backend = mm_state.backends[model_name]

            sub_messages: list[dict[str, str]] = []
            # Carry over system prompt
            for msg in messages:
                if msg.get("role") == "system":
                    sub_messages.append(msg)
                    break
            sub_messages.append({"role": "user", "content": task.text})

            grammar_str, is_json = _resolve_grammar(body, mm_state.default_json_mode)
            uses_grammar_natively = isinstance(backend, LlamaCppBackend)

            req_messages = list(sub_messages)
            if is_json and not uses_grammar_natively:
                rf = body.response_format or {}
                schema_for_prompt = None
                if rf.get("type") == "json_schema":
                    raw = rf.get("json_schema") or rf.get("schema")
                    schema_for_prompt = raw.get("schema", raw) if raw else None
                req_messages = _inject_json_system_prompt(req_messages, schema_for_prompt)

            gen_req = GenerationRequest(
                model=model_name,
                messages=req_messages,
                max_tokens=body.max_tokens,
                temperature=body.temperature,
                top_p=body.top_p,
                stream=False,
                grammar=grammar_str if uses_grammar_natively else None,
                json_mode=is_json,
            )

            mm_state.request_count += 1
            mm_state.routed_counts[model_name] = mm_state.routed_counts.get(model_name, 0) + 1

            try:
                text, _metrics = backend.generate(gen_req)
            except Exception as exc:
                logger.warning(
                    "Sub-task %d failed on model %s: %s",
                    task_idx,
                    model_name,
                    exc,
                )
                text = f"[Error processing sub-task {task_idx + 1}]"

            return SubTaskResult(
                task=task,
                response=text,
                model_used=model_name,
                tier=decision.tier,
            )

        # Build execution order respecting dependencies
        # Phase 1: execute independent tasks in parallel
        independent = [i for i, t in enumerate(decomp.tasks) if not t.depends_on]
        dependent = [i for i, t in enumerate(decomp.tasks) if t.depends_on]

        # Execute independent tasks concurrently
        if independent:
            results = await asyncio.gather(*[_execute_subtask(i) for i in independent])
            for result in results:
                sub_task_results[result.task.index] = result

        # Execute dependent tasks sequentially (in index order)
        for idx in sorted(dependent):
            # Wait for prerequisites (they should already be done)
            result = await _execute_subtask(idx)
            sub_task_results[result.task.index] = result

        # Merge results
        ordered_results = [sub_task_results[i] for i in range(len(decomp.tasks))]
        merged_text = decomposer_merger.merge(ordered_results)

        req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        response_data = {
            "id": req_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": ordered_results[0].model_used,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": merged_text},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
            },
        }

        resp = JSONResponse(content=response_data)
        resp.headers["X-Octomil-Decomposed"] = "true"
        resp.headers["X-Octomil-Subtasks"] = str(len(decomp.tasks))
        resp.headers["X-Octomil-Routed-Model"] = ordered_results[0].model_used
        resp.headers["X-Octomil-Tier"] = ordered_results[0].tier
        return resp

    return app


class _MultiModelStateAdapter:
    """Adapts MultiModelServerState for ``_stream_response`` compatibility.

    ``_stream_response`` expects a ``ServerState``-like object with
    ``.backend`` and ``.reporter``.  This adapter provides those.
    """

    def __init__(
        self,
        state: MultiModelServerState,
        backend: InferenceBackend,
        model_name: str,
    ) -> None:
        self.backend = backend
        self.reporter = state.reporter
        self.model_name = model_name


def run_multi_model_server(
    model_names: list[str],
    *,
    port: int = 8080,
    host: str = "0.0.0.0",
    api_key: Optional[str] = None,
    api_base: str = "https://api.octomil.com/api/v1",
    json_mode: bool = False,
    cache_size_mb: int = 2048,
    cache_enabled: bool = True,
    engine: Optional[str] = None,
    route_strategy: str = "complexity",
) -> None:
    """Start a multi-model inference server with query routing (blocking).

    Parameters
    ----------
    model_names:
        Ordered list of model names, smallest to largest.
    route_strategy:
        Routing strategy (``"complexity"`` is the only one currently).
    """
    import uvicorn

    app = create_multi_model_app(
        model_names,
        api_key=api_key,
        api_base=api_base,
        json_mode=json_mode,
        cache_size_mb=cache_size_mb,
        cache_enabled=cache_enabled,
        engine=engine,
        route_strategy=route_strategy,
    )
    uvicorn.run(app, host=host, port=port, log_level="info")
