"""MLX runtime: spawns llm_api.worker.mlx_worker in the server's own Python environment."""

from __future__ import annotations

import sys

import httpx

from .base import RuntimeAdapter, WorkerHandle


class MLXAdapter(RuntimeAdapter):
    name = "mlx"

    def available(self) -> bool:
        try:
            import mlx.core  # noqa: F401
            import mlx_lm  # noqa: F401
            return True
        except Exception:
            return False

    def build_command(self, model: dict, port: int, context: int) -> list[str]:
        py = self.settings.worker_python or sys.executable
        task = "embedding" if model.get("embedding") else "reranking" if model.get("reranker") else "text"
        cmd = [py, "-m", "llm_api.worker.mlx_worker", "--model-path", model["path"], "--model-id", model["id"],
               "--port", str(port), "--max-context", str(context),
               "--default-max-tokens", str(self.settings.default_max_tokens), "--task", task,
               "--generation-timeout", str(self.settings.generation_timeout_seconds)]
        if model.get("vision"):
            cmd.append("--vision")
        overrides = model.get("overrides") or {}
        if overrides.get("kv_bits"):
            cmd += ["--kv-bits", str(overrides["kv_bits"])]
        return cmd

    async def warmup(self, handle: WorkerHandle, client: httpx.AsyncClient, model: dict) -> dict:
        r = await client.post(f"{handle.base_url}/warmup", timeout=600)
        r.raise_for_status()
        return r.json()

    async def memory_gb(self, handle: WorkerHandle, client: httpx.AsyncClient) -> float | None:
        try:
            r = await client.get(f"{handle.base_url}/health", timeout=3)
            m = r.json().get("memory") or {}
            return (m.get("active_gb") or 0) + (m.get("cache_gb") or 0) or None
        except Exception:
            return None

    async def is_ready(self, handle: WorkerHandle, client: httpx.AsyncClient) -> tuple[str, str | None]:
        try:
            r = await client.get(f"{handle.base_url}/health", timeout=3)
        except Exception:
            return "loading", None
        if r.status_code != 200:
            return "loading", None
        d = r.json()
        return d.get("status", "loading"), d.get("error")