"""
MLX server and model management for the Advanced (Fine-Tuning) path.

Parallel to app_llmster.py — manages:
  - Isolated MLX venv (~/.mugi/mlx-venv/) for the mlx-vlm server, training, downloads
  - Separate utility MLX venv (~/.mugi/mlx-utility-venv/) for embedded background model (mlx-vlm; Prism MLX fork only needed for optional 1-bit Bonsai)
  - mlx-vlm server lifecycle (python -m mlx_vlm.server; start/stop/health on port 8000)
  - Model discovery (LM Studio MLX models + HuggingFace cache)
  - Adapter management (~/.mugi/mlx/adapters/)
  - Recommended mlx-community models

Does not import app_llmster or mugi_backend; safe for use from either.
"""
from __future__ import annotations

import json
import logging
import os
import re
import shlex
import shutil
import signal
import socket
import subprocess
import tempfile
import sys
import time
import traceback
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from threading import Lock

logger = logging.getLogger(__name__)

try:
    import requests
except ImportError:
    requests = None

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
MUGI_HOME = os.path.expanduser("~/.mugi")
MLX_VENV_DIR = os.path.join(MUGI_HOME, "mlx-venv")
# In-process utility model (utility_model.py): Prism MLX fork for 1-bit weights; keeps main mlx-venv on upstream mlx.
MLX_UTILITY_VENV_DIR = os.path.join(MUGI_HOME, "mlx-utility-venv")
MLX_DATA_DIR = os.path.join(MUGI_HOME, "mlx")
MLX_CANONICAL_MODELS_DIR = os.path.join(MLX_DATA_DIR, "models")
MLX_ADAPTERS_DIR = os.path.join(MLX_DATA_DIR, "adapters")
MLX_TRAINING_TMP_DIR = os.path.join(MLX_DATA_DIR, ".training_tmp")
MLX_VERSION_FILE = os.path.join(MLX_VENV_DIR, ".mlx_version")
APP_PREFS_PATH = os.path.join(MUGI_HOME, "app_prefs.json")
MLX_RUNTIME_STATE_FILE = os.path.join(MLX_DATA_DIR, "runtime_state.json")

# Bump when bundled ``_MLX_*_PATCHES`` tuples change so existing subprocesses reload patched site-packages.
# Gen 3: mlx-openai-server -> mlx-vlm server migration; all site-package patches retired (forces one restart).
# Gen 4: gemma4 tool parser accepts pythonic call:name(arg="v") syntax (QAT checkpoints drift from
#        the brace syntax; parser rejection silently dropped the whole completion -> empty replies).
_MLX_PATCH_GENERATION = 4
_MUGI_PATCH_GEN_FILE = os.path.join(MLX_DATA_DIR, ".mugi_patch_generation")

MLX_SERVER_PORT = 8000
MLX_SERVER_URL = f"http://127.0.0.1:{MLX_SERVER_PORT}"

# LM Studio model directories to scan for MLX-format models
_LMSTUDIO_MODEL_DIRS = [
    os.path.join(MUGI_HOME, ".lmstudio", "models"),
    os.path.expanduser("~/.lmstudio/models"),
]

# ---------------------------------------------------------------------------
# Recommended MLX models (mlx-community HuggingFace repos)
# ---------------------------------------------------------------------------
# Recommended chat = omni/mlx-vlm (multimodal). mlx-lm text-only launch is advanced/experimental.
MLX_RECOMMENDED: list[dict] = [
    # Primary recommendations by RAM tier (vision + native tools). Tiers are sized for the
    # FULL runtime footprint, not just weights: the resident model + the APC KV pool (which
    # pre-allocates ~64k tokens of KV up front — several GB) + the mlx-embeddings model +
    # macOS itself must all fit inside the GPU working set (~75% of unified RAM). Sizing by
    # weights alone over-subscribed low-RAM machines — a ~13 GB 8-bit 12B on a 24 GB Mac
    # thrashed to a generation stall (loads + serves /v1/models fine, then stalls the moment
    # real KV/generation pressure hits). Rule of thumb: usable weight budget is roughly
    # (RAM − ~16 GB reserved for OS + KV pool + embeddings).
    #   <= 16 GB : Gemma 4 E2B QAT (8-bit, ~6 GB)
    #   17–24 GB : Gemma 4 26B-A4B QAT (4-bit, ~7 GB)
    #   25–32 GB : Gemma 4 26B-A4B QAT (6-bit, ~10 GB)
    #   33–48 GB : Gemma 4 12B QAT (mxfp8, ~13 GB)
    #   > 48 GB  : Gemma 4 26B-A4B QAT (8-bit, ~29 GB)
    #
    # tool_call_parser / reasoning_parser: passed to mlx-openai-server so it
    # emits native OpenAI delta.tool_calls instead of dumping raw XML/channel
    # markup into content.  See _detect_tool_call_parsers() for heuristic
    # fallback when the model is not in this list.
    {
        "repo": "mlx-community/gemma-4-26B-A4B-it-qat-8bit",
        "display": "Gemma 4 26B-A4B MoE QAT (8-bit)",
        "size": "~29 GB",
        "ram": "More than 48 GB unified memory",
        "caps": ["tools", "vision"],
        "recommended_for_min_ram_gb": 49,
        "tool_call_parser": "gemma4",
    },
    {
        "repo": "mlx-community/gemma-4-12B-it-qat-mxfp8",
        "display": "Gemma 4 12B QAT (mxfp8)",
        "size": "~13 GB",
        "ram": "33–48 GB unified memory",
        "caps": ["tools", "vision"],
        "recommended_for_min_ram_gb": 33,
        "recommended_for_max_ram_gb": 48,
        "tool_call_parser": "gemma4",
    },
    {
        # Kept as a manual option but no longer an auto-recommended RAM tier: ~11 GB of
        # weights plus the APC KV pool over-subscribes the machines it used to target.
        "repo": "mlx-community/gemma-4-12B-it-qat-4bit",
        "display": "Gemma 4 12B QAT (4-bit)",
        "size": "~11 GB",
        "ram": "33 GB or more (manual option)",
        "caps": ["tools", "vision"],
        "tool_call_parser": "gemma4",
    },
    # Legacy / additional vision + tool-capable options (not RAM-tier recommended)
    {"repo": "mlx-community/Qwen3.6-27B-4bit", "display": "Qwen 3.6 27B (4-bit)", "size": "~18 GB", "ram": "24 GB or more", "caps": ["tools", "vision"], "revision": os.environ.get("MUGI_MLX_REV_QWEN36_27B", ""), "tool_call_parser": "qwen3_coder", "reasoning_parser": "qwen3_5"},
    {"repo": "mlx-community/Qwen3.6-27B-6bit", "display": "Qwen 3.6 27B (6-bit)", "size": "~24 GB", "ram": "32 GB or more", "caps": ["tools", "vision"], "tool_call_parser": "qwen3_coder", "reasoning_parser": "qwen3_5"},
    {"repo": "mlx-community/Qwen3.6-35B-A3B-6bit", "display": "Qwen 3.6 35B-A3B (6-bit)", "size": "~30 GB", "ram": "48 GB or more", "caps": ["tools", "vision"], "revision": os.environ.get("MUGI_MLX_REV_QWEN36_35B", ""), "tool_call_parser": "qwen3_coder", "reasoning_parser": "qwen3_5"},
    {"repo": "mlx-community/Qwen2.5-VL-7B-Instruct-4bit", "display": "Qwen 2.5 VL 7B (4-bit)", "size": "~5 GB", "ram": "6 GB", "caps": ["tools", "vision"], "tool_call_parser": "qwen3"},
    {"repo": "mlx-community/GLM-4.7-Flash-4bit", "display": "GLM-4.7 Flash (4-bit)", "size": "~6 GB", "ram": "8 GB", "caps": ["tools", "vision"], "tool_call_parser": "glm4_moe"},
    {"repo": "mlx-community/Mistral-Small-3.1-24B-Instruct-2503-4bit", "display": "Mistral Small 3.1 24B (4-bit)", "size": "~14 GB", "ram": "16 GB", "caps": ["tools", "vision"]},
    # Gemma 4 QAT — multimodal (vision). mlx-vlm 1.7.1+ exposes the real Gemma 4 parser.
    {"repo": "mlx-community/gemma-4-31B-it-qat-4bit", "display": "Gemma 4 31B QAT (4-bit)", "size": "~16 GB", "ram": "18 GB", "caps": ["tools", "vision"], "tool_call_parser": "gemma4"},
    {"repo": "mlx-community/gemma-4-31B-it-qat-mxfp8", "display": "Gemma 4 31B QAT (mxfp8)", "size": "~29 GB", "ram": "32 GB", "caps": ["tools", "vision"], "tool_call_parser": "gemma4"},
    {"repo": "mlx-community/gemma-4-26B-A4B-it-qat-6bit", "display": "Gemma 4 26B-A4B MoE QAT (6-bit)", "size": "~10 GB", "ram": "25–32 GB unified memory", "caps": ["tools", "vision"], "recommended_for_min_ram_gb": 25, "recommended_for_max_ram_gb": 32, "tool_call_parser": "gemma4"},
    {"repo": "mlx-community/gemma-4-26B-A4B-it-qat-4bit", "display": "Gemma 4 26B-A4B MoE QAT (4-bit)", "size": "~7 GB", "ram": "17–24 GB unified memory", "caps": ["tools", "vision"], "recommended_for_min_ram_gb": 17, "recommended_for_max_ram_gb": 24, "tool_call_parser": "gemma4"},
    {"repo": "mlx-community/gemma-4-E2B-it-qat-8bit", "display": "Gemma 4 E2B QAT (8-bit)", "size": "~6 GB", "ram": "16 GB or less", "caps": ["tools", "vision"], "recommended_for_max_ram_gb": 16, "tool_call_parser": "gemma4"},
]

# ---------------------------------------------------------------------------
# Venv management
# ---------------------------------------------------------------------------

# Fallback install list used ONLY when the bundled requirements files are missing.
# MUST stay in sync with the mlx-* direct deps in requirements-mlx-pinned.txt — a drift
# guard test (tests/test_mlx_install_dependencies.py) enforces this. mlx-embeddings is
# load-bearing for knowledge/session semantic search; omitting it ships a broken stack.
_MLX_PACKAGES = ["mlx", "mlx-metal", "mlx-lm[train]", "mlx-vlm", "mlx-embeddings>=0.0.5,<0.2"]
_MLX_STACK_PACKAGE_NAMES = ["mlx", "mlx-metal", "mlx-lm", "mlx-vlm"]
_MLX_TARGETED_UPGRADE_PACKAGES = {"mlx", "mlx-metal", "mlx-lm", "mlx-vlm"}
# Order for pre-release upgrades: pin mlx first, then the rest (matches requirements-mlx-pinned.txt).
_MLX_PRE_UPGRADE_PACKAGES = ["mlx", "mlx-metal", "mlx-lm[train]", "mlx-vlm"]
_MLX_REQUIREMENTS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "requirements-mlx.txt")
_MLX_PINNED_REQUIREMENTS_FILE = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "requirements-mlx-pinned.txt"
)
# Pip installs for mlx / mlx-lm / mlx-vlm often exceed 10 minutes on first install or reinstall.
_MLX_PIP_STACK_TIMEOUT_SEC = max(120, int(os.environ.get("MUGI_MLX_PIP_TIMEOUT", "3600")))
_UTILITY_MLX_PACKAGES = ["mlx-vlm"]
_PRISM_MLX_REQUIREMENT = "mlx @ git+https://github.com/PrismML-Eng/mlx.git@prism"
_PRISM_MLX_BUILD_LOG = os.path.join(MUGI_HOME, "logs", "prism_mlx_build.log")

# Written when `pip install` of Prism MLX succeeds (no public wheels — builds from source).
_UTILITY_PRISM_MARKER = os.path.join(MLX_UTILITY_VENV_DIR, ".mugi_prism_mlx_ok")

_PRISM_MLX_BUILD_HINT = (
    "Prism MLX has no prebuilt pip wheels; pip compiles the fork from source. On macOS, 1-bit Bonsai "
    "requires full Xcode—not just Command Line Tools—so the Metal compiler (`metal`) exists: "
    "install Xcode from the App Store, open it once, run `xcodebuild -downloadComponent MetalToolchain`, "
    "and `sudo xcode-select -s /Applications/Xcode.app/Contents/Developer`. "
    "CMake must be on PATH (e.g. `brew install cmake`). "
    "Then retry Settings → Advanced → Upgrade utility MLX, or run:\n  "
)


def utility_mlx_prism_ready() -> bool:
    """True when Prism MLX was installed successfully into the utility venv (1-bit Bonsai)."""
    return os.path.isfile(_UTILITY_PRISM_MARKER)


def _skip_prism_mlx_install() -> bool:
    return os.environ.get("MUGI_SKIP_PRISM_MLX", "").strip().lower() in ("1", "true", "yes")


def utility_model_requires_prism_mlx(model_repo: str) -> bool:
    """Return True if the given utility model repo needs the Prism MLX fork.

    Prism MLX is only required for the 1-bit Bonsai weights (and any other
    prism-ml/* fork builds). Stock mlx-community quantized models load via
    mlx-vlm alone, so we can skip the expensive Prism build for them.
    """
    ref = str(model_repo or "").strip().lower()
    if not ref:
        return False
    return "prism-ml/" in ref or "bonsai" in ref


def _utility_model_repo_from_config() -> str:
    """Best-effort lookup of the configured utility model repo.

    Imported lazily to avoid a circular import with core.config.
    """
    try:
        from core.config import get_utility_config  # local import to avoid cycles
        cfg = get_utility_config() or {}
        return str(cfg.get("model_repo") or "").strip()
    except Exception:
        return ""


def _mark_utility_prism_ok() -> None:
    try:
        os.makedirs(os.path.dirname(_UTILITY_PRISM_MARKER), exist_ok=True)
        with open(_UTILITY_PRISM_MARKER, "w", encoding="utf-8") as f:
            f.write("ok\n")
    except Exception:
        pass


def _clear_utility_prism_marker() -> None:
    try:
        os.remove(_UTILITY_PRISM_MARKER)
    except Exception:
        pass


def _augment_path_for_build_tools(env: dict[str, str]) -> dict[str, str]:
    """Add common toolchain paths (GUI app launches often miss Homebrew PATH)."""
    current = env.get("PATH", "")
    parts = [p for p in current.split(os.pathsep) if p]
    for p in ("/opt/homebrew/bin", "/usr/local/bin", "/usr/bin", "/bin"):
        if p not in parts:
            parts.append(p)
    env["PATH"] = os.pathsep.join(parts)
    return env


def _write_prism_build_log(stderr: str, stdout: str, *, phase: str = "pip") -> None:
    """Persist Prism MLX diagnostics (preflight, pip, or unexpected errors) for debugging."""
    try:
        os.makedirs(os.path.dirname(_PRISM_MLX_BUILD_LOG), exist_ok=True)
        with open(_PRISM_MLX_BUILD_LOG, "w", encoding="utf-8") as f:
            f.write(f"=== phase={phase} | ts={datetime.now().isoformat(timespec='seconds')} ===\n\n")
            f.write("=== STDERR / diagnostic text ===\n")
            f.write((stderr or "").strip())
            f.write("\n\n=== STDOUT / extra context ===\n")
            f.write((stdout or "").strip())
            f.write("\n")
    except Exception:
        pass


def _prism_log_hint() -> str:
    return f"Full diagnostic log: {_PRISM_MLX_BUILD_LOG}"


def _tail_for_error(text: str, max_chars: int = 6000) -> str:
    s = (text or "").strip()
    if not s:
        return ""
    return s[-max_chars:]


def _metal_toolchain_available(env: dict[str, str]) -> tuple[bool, str, str]:
    """Return (ok, user_message, xcrun_detail_for_log)."""
    if sys.platform != "darwin":
        return True, "", ""
    try:
        r = subprocess.run(
            ["xcrun", "--find", "metal"],
            capture_output=True,
            text=True,
            timeout=15,
            env=env,
        )
        detail = f"xcrun --find metal exit={r.returncode}\nstdout:\n{r.stdout or ''}\nstderr:\n{r.stderr or ''}"
        if r.returncode == 0 and (r.stdout or "").strip():
            return True, "", ""
        xcrun_err = (r.stderr or "").strip()
        xcrun_out = (r.stdout or "").strip()
        snippet = (xcrun_err or xcrun_out)[:500]
        lead = (
            "The Metal compiler (`metal`) is not available. Prism MLX builds need Apple's Metal "
            "toolchain, which requires full Xcode installed—Command Line Tools alone are usually not enough.\n\n"
        )
        if snippet:
            lead += f"xcrun said:\n  {snippet}\n\n"
        user = (
            f"{lead}"
            "Install Xcode from the App Store and open it once. Then in Terminal:\n"
            "  xcodebuild -downloadComponent MetalToolchain\n"
            "  sudo xcode-select -s /Applications/Xcode.app/Contents/Developer\n"
            "Quit and reopen Mugi, then tap Upgrade utility MLX again."
        )
        return False, user, detail
    except Exception as exc:
        return False, f"Could not verify Metal toolchain via xcrun: {exc}", traceback.format_exc()


def _install_prism_mlx_fork(pip_executable: str) -> tuple[bool, str]:
    """Overlay PrismML mlx (1-bit kernels) in the given venv (source build)."""
    env = {**os.environ, "HF_HUB_ENABLE_HF_TRANSFER": os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "1")}
    env = _augment_path_for_build_tools(env)
    py_executable = os.path.join(os.path.dirname(pip_executable), "python")
    py_bin_dir = os.path.dirname(py_executable)
    if os.path.isfile(py_executable):
        # Force CMake/scikit-build to compile Python bindings for this exact venv.
        existing_args = env.get("CMAKE_ARGS", "").strip()
        forced_args = (
            f"-DPython_EXECUTABLE={py_executable} "
            f"-DPython3_EXECUTABLE={py_executable} "
            f"-DPython_ROOT_DIR={os.path.dirname(py_bin_dir)}"
        )
        env["CMAKE_ARGS"] = (existing_args + " " + forced_args).strip()
        env["PYTHON_EXECUTABLE"] = py_executable
        env["Python_EXECUTABLE"] = py_executable
        env["Python3_EXECUTABLE"] = py_executable
        # Prefer the venv bin ahead of Homebrew/global bins.
        env["PATH"] = os.pathsep.join([py_bin_dir, env.get("PATH", "")])
    cmake_path = shutil.which("cmake", path=env.get("PATH", ""))
    if not cmake_path:
        path_ctx = env.get("PATH", "")
        user = (
            "Could not find `cmake` in PATH for the backend process. "
            "If this is a GUI launch, make sure Homebrew paths are visible to the app process "
            f"(current PATH={path_ctx})."
        )
        _write_prism_build_log(
            user,
            f"pip_executable={pip_executable}\npython_executable={py_executable if os.path.isfile(py_executable) else 'missing'}",
            phase="preflight_cmake",
        )
        return False, f"{user}\n\n{_prism_log_hint()}"
    metal_ok, metal_msg, metal_detail = _metal_toolchain_available(env)
    if not metal_ok:
        _write_prism_build_log(metal_msg, metal_detail, phase="preflight_metal")
        return False, f"{metal_msg}\n\n{_prism_log_hint()}"
    try:
        r = subprocess.run(
            [pip_executable, "install", "--upgrade", "--no-cache-dir", "-v", _PRISM_MLX_REQUIREMENT],
            check=True,
            capture_output=True,
            text=True,
            timeout=3600,
            env=env,
        )
        _write_prism_build_log("", r.stdout or "", phase="pip_success")
        return True, ""
    except subprocess.CalledProcessError as e:
        stderr = e.stderr or ""
        stdout = e.stdout or ""
        _write_prism_build_log(stderr, stdout, phase="pip_failed")
        combined = "\n".join(
            part for part in (
                _tail_for_error(stderr),
                _tail_for_error(stdout),
            ) if part
        )
        if "cannot execute tool 'metal'" in combined:
            combined += (
                "\n\nDetected missing Metal toolchain. Install it once with:\n"
                "  xcodebuild -downloadComponent MetalToolchain"
            )
        return False, (
            f"{combined}\n\n"
            f"{_prism_log_hint()}"
        ).strip()
    except Exception as e:
        tb = traceback.format_exc()
        _write_prism_build_log(str(e), tb, phase="unexpected_error")
        return False, f"{e}\n\n{_prism_log_hint()}"


def _try_prism_mlx_for_utility_venv() -> tuple[bool, str]:
    """Install Prism MLX if needed; on success write marker. Returns (prism_ok, pip_error_or_empty)."""
    if _skip_prism_mlx_install():
        return True, ""
    if utility_mlx_prism_ready():
        return True, ""
    # Remove any broken prior build (e.g. cp314 extension in a cp312 venv) before retry.
    try:
        subprocess.run(
            [_utility_mlx_pip(), "uninstall", "-y", "mlx"],
            check=False,
            capture_output=True,
            text=True,
            timeout=120,
        )
    except Exception:
        pass
    ok, err = _install_prism_mlx_fork(_utility_mlx_pip())
    if ok:
        _mark_utility_prism_ok()
        return True, ""
    return False, err


def mlx_venv_exists() -> bool:
    """Return True if the MLX venv exists and has a pip binary."""
    pip = os.path.join(MLX_VENV_DIR, "bin", "pip")
    return os.path.isfile(pip) and os.access(pip, os.X_OK)


def _mlx_pip() -> str:
    return os.path.join(MLX_VENV_DIR, "bin", "pip")


def _mlx_python() -> str:
    return os.path.join(MLX_VENV_DIR, "bin", "python")


def utility_mlx_venv_exists() -> bool:
    """True if the dedicated utility-model venv is present."""
    pip = os.path.join(MLX_UTILITY_VENV_DIR, "bin", "pip")
    return os.path.isfile(pip) and os.access(pip, os.X_OK)


def _utility_mlx_pip() -> str:
    return os.path.join(MLX_UTILITY_VENV_DIR, "bin", "pip")


def utility_mlx_python() -> str:
    """Python interpreter for the utility MLX venv (snapshot_download, mlx_vlm)."""
    return os.path.join(MLX_UTILITY_VENV_DIR, "bin", "python")


def ensure_utility_mlx_venv() -> tuple[bool, str]:
    """Create ~/.mugi/mlx-utility-venv with mlx-vlm.

    Prism MLX (required only by the 1-bit Bonsai weights) is installed lazily:
    it's skipped unless the configured utility model actually needs it, or the
    Prism marker is already present (user opted in earlier). Set
    MUGI_SKIP_PRISM_MLX=1 to force-skip Prism even when Bonsai is selected.
    """
    if not utility_mlx_venv_exists():
        os.makedirs(MLX_UTILITY_VENV_DIR, exist_ok=True)
        try:
            subprocess.run(
                [sys.executable, "-m", "venv", MLX_UTILITY_VENV_DIR],
                check=True, capture_output=True, text=True, timeout=60,
            )
        except Exception as e:
            return False, f"Failed to create utility MLX venv: {e}"
        try:
            subprocess.run(
                [_utility_mlx_pip(), "install", "--upgrade", "pip"],
                check=True, capture_output=True, text=True, timeout=120,
            )
            subprocess.run(
                [_utility_mlx_pip(), "install"] + _UTILITY_MLX_PACKAGES,
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        except subprocess.CalledProcessError as e:
            return False, f"Utility MLX pip install failed: {e.stderr or e.stdout or str(e)}"
        except Exception as e:
            return False, f"Utility MLX pip install failed: {e}"

    # Skip the (expensive, Xcode-dependent) Prism MLX build unless the user
    # actually needs it for the currently configured utility model. If Prism
    # was successfully installed previously we keep it in place regardless.
    needs_prism = utility_model_requires_prism_mlx(_utility_model_repo_from_config())
    if not needs_prism and not utility_mlx_prism_ready():
        return True, "Utility MLX environment ready."

    prism_ok, prism_err = _try_prism_mlx_for_utility_venv()
    if prism_ok:
        if _skip_prism_mlx_install():
            return True, (
                "Utility MLX environment ready. Prism install skipped (MUGI_SKIP_PRISM_MLX); "
                "1-bit Bonsai will not load until Prism MLX is installed."
            )
        return True, "Utility MLX environment ready."
    manual = f'{_PRISM_MLX_BUILD_HINT}{_utility_mlx_pip()} install --upgrade "{_PRISM_MLX_REQUIREMENT}"'
    return True, (
        "Utility MLX venv is ready for standard quantized models. "
        "1-bit Bonsai still needs a successful Prism MLX build:\n\n"
        f"{prism_err}\n\n{manual}"
    )


def upgrade_utility_mlx_venv() -> tuple[bool, str]:
    """Upgrade mlx-vlm and re-apply Prism MLX in the utility venv."""
    if not utility_mlx_venv_exists():
        return False, "Utility MLX venv not installed."
    errors: list[str] = []
    for pkg in _UTILITY_MLX_PACKAGES:
        try:
            subprocess.run(
                [_utility_mlx_pip(), "install", "--upgrade", pkg],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        except subprocess.CalledProcessError as e:
            errors.append(f"{pkg}: {e.stderr or e.stdout or str(e)}")
        except Exception as e:
            errors.append(f"{pkg}: {e}")
    if errors:
        return False, "Utility MLX upgrade errors:\n" + "\n".join(errors)

    had_prism = utility_mlx_prism_ready()
    needs_prism = utility_model_requires_prism_mlx(_utility_model_repo_from_config())
    _clear_utility_prism_marker()
    if not needs_prism and not had_prism:
        return True, (
            "mlx-vlm upgraded. Prism MLX not installed (not needed for the current utility model)."
        )

    prism_ok, prism_err = _try_prism_mlx_for_utility_venv()
    if not prism_ok and not _skip_prism_mlx_install():
        manual = f'{_PRISM_MLX_BUILD_HINT}{_utility_mlx_pip()} install --upgrade "{_PRISM_MLX_REQUIREMENT}"'
        return True, (
            "mlx-vlm upgraded. Prism MLX build failed (1-bit Bonsai unavailable until fixed):\n\n"
            f"{prism_err}\n\n{manual}"
        )
    if _skip_prism_mlx_install():
        return True, "mlx-vlm upgraded. Prism skipped (MUGI_SKIP_PRISM_MLX)."
    return True, "Utility MLX packages upgraded."


def delete_utility_mlx_venv() -> tuple[bool, str]:
    """Remove only the utility MLX venv (not models under ~/.mugi/mlx/)."""
    if os.path.isdir(MLX_UTILITY_VENV_DIR):
        try:
            shutil.rmtree(MLX_UTILITY_VENV_DIR)
        except Exception as e:
            return False, f"Could not remove utility MLX venv: {e}"
    return True, "Utility MLX environment removed."


# ---------------------------------------------------------------------------
# Post-install patches for mlx-vlm
# ---------------------------------------------------------------------------

# Gen 4 — gemma4 tool parser: accept pythonic call syntax.
#
# Gemma-4 QAT checkpoints sometimes emit `call:name(arg="value")` (parens, plain
# quotes) instead of the expected `call:name{arg:<|"|>value<|"|>}` brace syntax.
# Upstream parse_tool_call raises on the paren form, and process_tool_calls then
# DROPS the whole completion ("Invalid tool call" in server.log) — the client
# receives an empty 200 stream. These two patches teach the parser the paren
# form so the call is emitted as a native tool_calls chunk instead of vanishing.
#
# (Historical note: the mlx-vlm 0.4.4 thread-local-stream patches were retired in
# Gen 3; restore from git history pre-migration if ever needed.)
_GEMMA4_PAREN_HELPER = '''
def _parse_paren_call_args(args_str):
    """Mugi patch: parse pythonic call arguments (key="value", n=3) into a dict.

    Some gemma-4 QAT checkpoints drift from the brace syntax and emit
    call:name(key="value") instead of call:name{key:<|"|>value<|"|>}.
    JSON-style bare literals (true/false/null) are mapped to Python values.
    """
    import ast

    args_str = args_str.strip()
    if not args_str:
        return {}
    tree = ast.parse("_f({})".format(args_str), mode="eval")
    _NAME_LITERALS = {"true": True, "false": False, "null": None, "none": None}
    out = {}
    for kw in tree.body.keywords:
        if kw.arg is None:
            continue
        value = kw.value
        if isinstance(value, ast.Name) and value.id.lower() in _NAME_LITERALS:
            out[kw.arg] = _NAME_LITERALS[value.id.lower()]
        else:
            out[kw.arg] = ast.literal_eval(value)
    return out


def parse_tool_call(text, tools=None):
    # Mugi patch applied. Try recursive-group regex first for well-formed calls'''

# NOTE: each patch's old_text must NOT survive in the patched file, or
# _apply_mlx_venv_patches would re-apply it on every start (it replaces the
# first occurrence whenever old_text is found). Both anchors below vanish
# after patching: the first because its comment line is rewritten, the second
# because the raise is no longer adjacent to `if not m:`.
_MLX_VLM_PATCHES: list[tuple[str, str, str]] = [
    (
        "mlx_vlm/tool_parsers/gemma4.py",
        '''def parse_tool_call(text, tools=None):
    # Try recursive-group regex first for well-formed calls''',
        _GEMMA4_PAREN_HELPER,
    ),
    (
        "mlx_vlm/tool_parsers/gemma4.py",
        '''        m = re.search(r"call:([\\w-]+)\\{", text)
        if not m:
            raise ValueError("No function call found in tool call text.")''',
        '''        m = re.search(r"call:([\\w-]+)\\{", text)
        if not m:
            # Mugi patch: tolerate pythonic paren syntax call:name(arg="v").
            paren = re.search(r"call:([\\w-]+)\\((.*)\\)\\s*$", text.strip(), re.DOTALL)
            if paren is not None:
                arguments = _parse_paren_call_args(paren.group(2))
                return dict(name=paren.group(1), arguments=json.dumps(arguments))
            raise ValueError("No function call found in tool call text.")''',
    ),
]

# mlx-openai-server is retired; its `app/*` package is no longer installed, so its
# thread-local-stream patches are dead. Kept as an empty list so _apply_mlx_venv_patches
# stays generic if patches are ever reintroduced.
_MLX_OAI_SERVER_PATCHES: list[tuple[str, str, str]] = []


def _apply_mlx_venv_patches() -> bool:
    """Apply workaround patches to installed mlx-vlm and mlx-openai-server files.

    Idempotent — skips each patch if the old text is absent (already patched
    or upstream changed).  Called after every pip install / upgrade.

    After writing a patched .py file, the corresponding __pycache__/*.pyc is
    deleted and the .py mtime is bumped so that Python never serves stale
    bytecode compiled from the pre-patch source.

    Returns:
        True if at least one site-packages file was modified.
    """
    sp = os.path.join(MLX_VENV_DIR, "lib")
    if not os.path.isdir(sp):
        return False
    try:
        pydir = next(
            d for d in sorted(os.listdir(sp))
            if d.startswith("python") and os.path.isdir(os.path.join(sp, d, "site-packages"))
        )
    except StopIteration:
        return False
    site = os.path.join(sp, pydir, "site-packages")

    import glob as _glob

    applied_any = False
    _all = _MLX_VLM_PATCHES + _MLX_OAI_SERVER_PATCHES
    for rel_path, old_text, new_text in _all:
        fpath = os.path.join(site, rel_path)
        if not os.path.isfile(fpath):
            continue
        try:
            src = open(fpath, "r", encoding="utf-8").read()
            if old_text not in src:
                continue
            src = src.replace(old_text, new_text, 1)
            with open(fpath, "w", encoding="utf-8") as f:
                f.write(src)

            # Invalidate .pyc bytecache so the server loads patched source.
            stem = os.path.splitext(os.path.basename(fpath))[0]
            cache_dir = os.path.join(os.path.dirname(fpath), "__pycache__")
            for pyc in _glob.glob(os.path.join(cache_dir, f"{stem}.*.pyc")):
                try:
                    os.remove(pyc)
                except OSError:
                    pass

            # Bump mtime by 2 s so even surviving .pyc is seen as stale.
            future = time.time() + 2
            os.utime(fpath, (future, future))

            applied_any = True
            logger.info("MLX venv patch applied: %s", rel_path)
        except Exception:
            logger.warning("MLX venv patch failed for %s", rel_path, exc_info=True)
    return applied_any


def mlx_prepare_venv_runtime() -> bool:
    """Apply bundled site-package patches and return True if mlx-openai-server should restart.

    Writes ``~/.mugi/mlx/.mugi_patch_generation`` when Mugi ships a new patch generation so an
    already-running server (same model fingerprint) still reloads — patching disk alone does not
    fix an old subprocess.
    """
    os.makedirs(MLX_DATA_DIR, exist_ok=True)
    try:
        prev = int(open(_MUGI_PATCH_GEN_FILE, encoding="utf-8").read().strip())
    except Exception:
        prev = 0
    generation_stale = prev != _MLX_PATCH_GENERATION
    patched_files = _apply_mlx_venv_patches()
    if generation_stale:
        try:
            with open(_MUGI_PATCH_GEN_FILE, "w", encoding="utf-8") as f:
                f.write(str(_MLX_PATCH_GENERATION))
        except Exception:
            logger.debug("Failed to write MLX patch generation marker", exc_info=True)
        return True
    return patched_files


def _mlx_install_requirements_path(*, prefer_pinned: bool = True) -> str | None:
    """Requirements file for venv create/reinstall (pinned stable stack by default)."""
    if prefer_pinned and os.path.isfile(_MLX_PINNED_REQUIREMENTS_FILE):
        return _MLX_PINNED_REQUIREMENTS_FILE
    if os.path.isfile(_MLX_REQUIREMENTS_FILE):
        return _MLX_REQUIREMENTS_FILE
    return None


def _mlx_stack_is_complete() -> bool:
    """True only if the MLX venv can actually import the full runtime + embeddings stack.

    ``mlx_venv_exists()`` only checks for ``bin/pip``; that is NOT enough — an interrupted
    or partial install (e.g. the user hit Retry mid-install) leaves a venv directory that
    looks present but cannot serve chat or build the vector index. Verify the real imports
    so a half-built environment gets repaired instead of trusted.
    """
    py = _mlx_python()
    if not os.path.isfile(py):
        return False
    try:
        proc = subprocess.run(
            [py, "-c", "import mlx, mlx_lm, mlx_vlm, mlx_embeddings"],
            capture_output=True, text=True, timeout=120,
        )
        return proc.returncode == 0
    except Exception:
        return False


def ensure_mlx_venv(*, prefer_pinned: bool = True) -> tuple[bool, str]:
    """Create the MLX venv and install packages. Returns (success, message).

    Default installs ``requirements-mlx-pinned.txt`` (known-good stack). Use
    ``prefer_pinned=False`` only when callers intentionally want floating deps.

    A pre-existing venv is reused only when its stack actually imports; an incomplete
    venv (partial/interrupted install) is repaired by reinstalling into it.
    """
    needs_create = not mlx_venv_exists()
    if not needs_create:
        if _mlx_stack_is_complete():
            return True, "MLX venv already exists."
        logger.warning("MLX venv present but stack incomplete; reinstalling MLX packages.")
    if needs_create:
        os.makedirs(MLX_VENV_DIR, exist_ok=True)
        try:
            subprocess.run(
                [sys.executable, "-m", "venv", MLX_VENV_DIR],
                check=True, capture_output=True, text=True, timeout=60,
            )
        except Exception as e:
            return False, f"Failed to create venv: {e}"
    req_path = _mlx_install_requirements_path(prefer_pinned=prefer_pinned)
    try:
        subprocess.run(
            [_mlx_pip(), "install", "--upgrade", "pip"],
            check=True, capture_output=True, text=True, timeout=120,
        )
        if req_path:
            subprocess.run(
                [_mlx_pip(), "install", "-r", req_path],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        else:
            subprocess.run(
                [_mlx_pip(), "install"] + _MLX_PACKAGES,
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
    except subprocess.CalledProcessError as e:
        return False, f"pip install failed: {e.stderr or e.stdout or str(e)}"
    except Exception as e:
        return False, f"pip install failed: {e}"
    _save_installed_version()
    mlx_prepare_venv_runtime()
    ver = get_installed_version() or "?"
    label = "pinned stable stack" if req_path == _MLX_PINNED_REQUIREMENTS_FILE else "MLX requirements"
    return True, f"MLX environment installed ({label}, mlx-lm {ver})."


def _pip_upgrade_mlx_core_and_vlm_follow_up() -> tuple[bool, str]:
    """Explicit second pass for wheels omitted from or skipped by batch `-r` resolution.

    Some resolver paths satisfy `mlx-vlm` transitively at an older pin even though PyPI has a
    newer release; forcing these two upgrades keeps Settings → Update stack aligned with the
    orange PyPI hints (mlx may already be current via mlx-lm—this is idempotent when satisfied).
    """
    errs: list[str] = []
    for pkg in ("mlx", "mlx-metal", "mlx-vlm"):
        try:
            subprocess.run(
                [_mlx_pip(), "install", "--upgrade", pkg],
                check=True,
                capture_output=True,
                text=True,
                timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        except subprocess.CalledProcessError as e:
            errs.append(f"{pkg}: {e.stderr or e.stdout or str(e)}")
        except Exception as e:
            errs.append(f"{pkg}: {e}")
    if errs:
        return False, "Follow-up mlx/mlx-vlm upgrade:\n" + "\n".join(errs)
    return True, ""


def _selected_mlx_model_and_context_from_prefs() -> tuple[str, int | None]:
    """Return (model_path, context_length) from app prefs for post-upgrade health checks."""
    prefs = _load_app_prefs_for_mlx()
    model_path = str((prefs or {}).get("mlx_model") or "").strip()
    ctx = None
    raw_ctx = (prefs or {}).get("llm_context_length")
    try:
        ctx_i = int(raw_ctx)
    except Exception:
        ctx_i = 0
    if _MLX_CONTEXT_LENGTH_MIN <= ctx_i <= _MLX_CONTEXT_LENGTH_MAX:
        ctx = ctx_i
    return model_path, ctx


def _mlx_chat_verify_timeout_seconds() -> float:
    raw = str(os.environ.get("MUGI_MLX_VERIFY_TIMEOUT") or "90").strip()
    try:
        return max(25.0, min(float(raw), 300.0))
    except ValueError:
        return 90.0


def _verify_local_mlx_chat_once(timeout_seconds: float | None = None) -> tuple[bool, str]:
    """Best-effort non-stream chat probe against the local mlx-vlm server."""
    timeout = float(timeout_seconds if timeout_seconds is not None else _mlx_chat_verify_timeout_seconds())
    base = f"{MLX_SERVER_URL}/v1"
    short_timeout = max(10.0, timeout / 3.0)
    try:
        import urllib.request

        model_id = ""
        # Prefer /health.loaded_model. mlx-vlm's /v1/models lists the entire HF cache, so
        # data[0].id is not necessarily the loaded model — chatting against the wrong id
        # would evict the loaded model and load another. /health reports the real one.
        try:
            with urllib.request.urlopen(MLX_SERVER_URL + "/health", timeout=short_timeout) as resp:
                health = json.loads(resp.read().decode("utf-8", errors="replace"))
            if isinstance(health, dict):
                model_id = str(health.get("loaded_model") or "").strip()
        except Exception:
            model_id = ""
        if not model_id:
            with urllib.request.urlopen(base + "/models", timeout=short_timeout) as resp:
                payload = json.loads(resp.read().decode("utf-8", errors="replace"))
            data = payload.get("data") if isinstance(payload, dict) else []
            if isinstance(data, list) and data:
                model_id = str((data[0] or {}).get("id") or "").strip()
        if not model_id:
            return False, "no model returned by /health or /v1/models"
        req = urllib.request.Request(
            base + "/chat/completions",
            data=json.dumps(
                {
                    "model": model_id,
                    "messages": [{"role": "user", "content": "Respond with: ok"}],
                    "max_tokens": 8,
                    "temperature": 0.0,
                    "stream": False,
                }
            ).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            _ = resp.read()
        return True, ""
    except Exception as e:
        return False, str(e)


def _restart_selected_model_and_verify_chat(reason_label: str) -> tuple[bool, str]:
    """Restart selected MLX model after stack mutation and ensure chat is responsive."""
    model_path, context_length = _selected_mlx_model_and_context_from_prefs()
    if not model_path:
        return True, f"{reason_label} No MLX model selected; run Load model before chatting."
    ok_start, msg_start = start_mlx_server(model_path, context_length=context_length)
    if not ok_start:
        return False, (
            f"{reason_label} Stack changed but selected model could not restart: {msg_start}"
        )
    # Large models often need tens of seconds after process start before first token.
    time.sleep(5.0)
    last_err = ""
    verify_timeout = _mlx_chat_verify_timeout_seconds()
    for attempt in range(4):
        ok_verify, err = _verify_local_mlx_chat_once(timeout_seconds=verify_timeout)
        if ok_verify:
            _run_and_persist_mlx_native_tool_probe()
            return True, f"{reason_label} Selected model restarted and chat verification passed."
        last_err = err
        if attempt < 3:
            time.sleep(5.0)
    return False, f"{reason_label} Selected model restarted but chat verify failed: {last_err}"


def clear_mlx_last_tool_probe() -> None:
    """Drop cached native-tool probe so callers do not trust stale results after a restart."""
    try:
        from utils.app_prefs_store import load_app_prefs_cached, save_app_prefs

        prefs = load_app_prefs_cached()
        if not isinstance(prefs, dict) or "mlx_last_tool_probe" not in prefs:
            return
        updated = dict(prefs)
        del updated["mlx_last_tool_probe"]
        save_app_prefs(updated)
    except Exception:
        logger.debug("Failed to clear mlx_last_tool_probe", exc_info=True)


def _run_and_persist_mlx_native_tool_probe() -> None:
    """Best-effort native tool probe after stack restart; failures are non-fatal."""
    try:
        from utils.llm_tool_call_probe import probe_streaming_native_tool_calls
        from utils.app_prefs_store import load_app_prefs_cached, save_app_prefs

        result = probe_streaming_native_tool_calls(
            MLX_SERVER_URL + "/v1",
            api_key="lm-studio",
            max_tokens=96,
            timeout=90.0,
        )
        prefs = load_app_prefs_cached()
        if not isinstance(prefs, dict):
            return
        updated = dict(prefs)
        updated["mlx_last_tool_probe"] = {
            "ts": int(time.time()),
            "ok": bool(result.get("ok")),
            "native_tool_calls": bool(result.get("native_tool_calls")),
            "native_tool_calls_ok": bool(result.get("native_tool_calls")),
            "xml_like_in_content": bool(result.get("xml_like_in_content")),
            "probe_status": "native_ok" if result.get("native_tool_calls") else str(result.get("error_code") or "no_native_tool_calls"),
            "error_code": result.get("error_code"),
            "message": str(result.get("message") or "")[:500],
            "model": str(result.get("model") or "")[:200],
        }
        save_app_prefs(updated)
    except Exception:
        logger.debug("Post-restart native tool probe failed", exc_info=True)


def _revert_verify_soft_success_message(ver: str, health_msg: str) -> str:
    """Pinned pip install succeeded; post-revert chat probe alone failed."""
    detail = health_msg
    if "chat verify failed:" in health_msg:
        detail = health_msg.split("chat verify failed:", 1)[-1].strip()
    return (
        f"Pinned stack installed (mlx-lm {ver}). The quick chat check timed out or failed "
        f"({detail or 'timed out'}) — common for large models right after a stack change. "
        "Use Load model in Settings → LLMs to restart the server, then try chat."
    )


def upgrade_mlx_venv() -> tuple[bool, str]:
    """Upgrade packages in the existing venv from requirements-mlx.txt (Settings → “Update to …”).

    pip enforces only what each package’s metadata requires (and may upgrade mlx-lm, mlx, etc.
    on different cadences from mlx-vlm or mlx-openai-server). It does *not* guarantee a
    combination that the whole stack has been integration-tested with—see requirements-mlx.txt
    header and Revert to stable (requirements-mlx-pinned.txt).

    Upgrades each pass separately when not using the requirements file to avoid pip
    backtracking on cross-package dependency conflicts.
    """
    if not mlx_venv_exists():
        return False, "MLX venv not installed."
    if os.path.isfile(_MLX_REQUIREMENTS_FILE):
        try:
            subprocess.run(
                [_mlx_pip(), "install", "--upgrade", "-r", _MLX_REQUIREMENTS_FILE],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
            ok_follow, follow_msg = _pip_upgrade_mlx_core_and_vlm_follow_up()
            if not ok_follow:
                return False, follow_msg
            _save_installed_version()
            mlx_prepare_venv_runtime()
            ok_health, health_msg = _restart_selected_model_and_verify_chat("MLX packages upgraded from requirements-mlx.txt.")
            if not ok_health:
                return False, health_msg
            return True, health_msg
        except subprocess.CalledProcessError as e:
            return False, f"Upgrade errors: {e.stderr or e.stdout or str(e)}"
        except Exception as e:
            return False, f"Upgrade errors: {e}"
    errors: list[str] = []
    for pkg in _MLX_PACKAGES:
        try:
            subprocess.run(
                [_mlx_pip(), "install", "--upgrade", pkg],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        except subprocess.CalledProcessError as e:
            errors.append(f"{pkg}: {e.stderr or e.stdout or str(e)}")
        except Exception as e:
            errors.append(f"{pkg}: {e}")
    _save_installed_version()
    mlx_prepare_venv_runtime()
    if errors:
        return False, "Upgrade errors:\n" + "\n".join(errors)
    return True, "MLX packages upgraded."


def upgrade_mlx_venv_pre_releases() -> tuple[bool, str]:
    """Upgrade the MLX venv allowing PyPI pre-releases (pip --pre).

    Uses requirements-mlx.txt when present; otherwise upgrades mlx then mlx-lm, mlx-vlm,
    mlx-openai-server in separate passes (same strategy as upgrade_mlx_venv). Pre-releases
    only apply when maintainers publish them to PyPI; there is no separate nightly index.
    """
    if not mlx_venv_exists():
        return False, "MLX venv not installed."
    try:
        subprocess.run(
            [_mlx_pip(), "install", "--upgrade", "pip"],
            check=True, capture_output=True, text=True, timeout=120,
        )
    except subprocess.CalledProcessError as e:
        return False, f"pip upgrade failed: {e.stderr or e.stdout or str(e)}"
    except Exception as e:
        return False, f"pip upgrade failed: {e}"

    if os.path.isfile(_MLX_REQUIREMENTS_FILE):
        try:
            subprocess.run(
                [_mlx_pip(), "install", "--upgrade", "--pre", "-r", _MLX_REQUIREMENTS_FILE],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
            ok_follow, follow_msg = _pip_upgrade_mlx_core_and_vlm_follow_up()
            if not ok_follow:
                return False, follow_msg
            _save_installed_version()
            mlx_prepare_venv_runtime()
            return True, "MLX packages upgraded from requirements-mlx.txt (pre-releases allowed)."
        except subprocess.CalledProcessError as e:
            return False, f"Pre-release upgrade errors: {e.stderr or e.stdout or str(e)}"
        except Exception as e:
            return False, f"Pre-release upgrade errors: {e}"

    errors: list[str] = []
    for pkg in _MLX_PRE_UPGRADE_PACKAGES:
        try:
            subprocess.run(
                [_mlx_pip(), "install", "--upgrade", "--pre", pkg],
                check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
            )
        except subprocess.CalledProcessError as e:
            errors.append(f"{pkg}: {e.stderr or e.stdout or str(e)}")
        except Exception as e:
            errors.append(f"{pkg}: {e}")
    _save_installed_version()
    mlx_prepare_venv_runtime()
    if errors:
        return False, "Pre-release upgrade errors:\n" + "\n".join(errors)
    return True, "MLX packages upgraded (pre-releases allowed)."


def get_pinned_mlx_stack_summary() -> str:
    """Human-readable list of the pinned 'stable' stack (from requirements-mlx-pinned.txt)."""
    if not os.path.isfile(_MLX_PINNED_REQUIREMENTS_FILE):
        return ""
    parts: list[str] = []
    try:
        with open(_MLX_PINNED_REQUIREMENTS_FILE, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                # e.g. mlx-lm[train]==0.31.1 -> mlx-lm 0.31.1
                req = line.split("#", 1)[0].strip()
                if "==" in req:
                    name, ver = req.split("==", 1)
                    name = name.split("[", 1)[0]
                    parts.append(f"{name} {ver}")
    except OSError:
        return ""
    return ", ".join(parts) if parts else ""


def revert_mlx_venv_to_pinned() -> tuple[bool, str]:
    """Install exact versions from requirements-mlx-pinned.txt (Settings → Revert to stable)."""
    if not mlx_venv_exists():
        return False, "MLX venv not installed."
    if not os.path.isfile(_MLX_PINNED_REQUIREMENTS_FILE):
        return False, "Pinned requirements missing (requirements-mlx-pinned.txt)."
    stop_mlx_server()
    try:
        subprocess.run(
            [_mlx_pip(), "install", "--upgrade", "pip"],
            check=True, capture_output=True, text=True, timeout=120,
        )
        subprocess.run(
            [_mlx_pip(), "install", "-r", _MLX_PINNED_REQUIREMENTS_FILE],
            check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
        )
    except subprocess.CalledProcessError as e:
        return False, f"Revert failed: {e.stderr or e.stdout or str(e)}"
    except Exception as e:
        return False, f"Revert failed: {e}"
    _save_installed_version()
    mlx_prepare_venv_runtime()
    ver = get_installed_version() or "?"
    ok_health, health_msg = _restart_selected_model_and_verify_chat(
        f"Reverted to pinned stack (mlx-lm {ver})."
    )
    if not ok_health:
        return True, _revert_verify_soft_success_message(ver, health_msg)
    return True, health_msg


def delete_mlx_venv(remove_data: bool = False) -> tuple[bool, str]:
    """Remove the MLX venv and optionally all MLX data (models, adapters)."""
    stop_mlx_server()
    if os.path.isdir(MLX_VENV_DIR):
        try:
            shutil.rmtree(MLX_VENV_DIR)
        except Exception as e:
            return False, f"Could not remove venv: {e}"
    if remove_data and os.path.isdir(MLX_DATA_DIR):
        try:
            shutil.rmtree(MLX_DATA_DIR)
        except Exception as e:
            return False, f"Venv removed but could not remove data: {e}"
    return True, "MLX environment removed."


def _save_installed_version():
    """Record installed mlx-lm version for update checking."""
    try:
        r = subprocess.run(
            [_mlx_pip(), "show", "mlx-lm"],
            capture_output=True, text=True, timeout=30,
        )
        for line in (r.stdout or "").splitlines():
            if line.lower().startswith("version:"):
                version = line.split(":", 1)[1].strip()
                os.makedirs(os.path.dirname(MLX_VERSION_FILE), exist_ok=True)
                with open(MLX_VERSION_FILE, "w") as f:
                    f.write(version)
                return
    except Exception:
        pass


def get_installed_version() -> str | None:
    """Return the installed mlx-lm version string, or None."""
    try:
        if os.path.isfile(MLX_VERSION_FILE):
            with open(MLX_VERSION_FILE, "r") as f:
                return f.read().strip() or None
    except Exception:
        pass
    return None


def get_installed_package_version(package: str) -> str | None:
    """Return installed version for an allowlisted MLX stack package."""
    pkg = (package or "").strip()
    if pkg not in _MLX_TARGETED_UPGRADE_PACKAGES:
        return None
    if not mlx_venv_exists():
        return None
    try:
        r = subprocess.run(
            [_mlx_pip(), "show", pkg],
            capture_output=True, text=True, timeout=30,
        )
        if r.returncode != 0:
            return None
        for line in (r.stdout or "").splitlines():
            if line.lower().startswith("version:"):
                return line.split(":", 1)[1].strip() or None
    except Exception:
        pass
    return None


def _pip_index_timeout_seconds(timeout_seconds: float | int | None, default: int = 30) -> int:
    try:
        timeout = int(float(timeout_seconds)) if timeout_seconds is not None else int(default)
    except (TypeError, ValueError):
        timeout = int(default)
    return max(1, min(timeout, 120))


def _pip_index_latest_version(stdout: str) -> str | None:
    """Parse ``pip index versions`` stdout into the newest release string."""
    text = stdout or ""
    m = re.search(r"(?im)^\s*LATEST:\s*(\S+)", text)
    if m:
        return (m.group(1) or "").strip() or None
    first = next((ln.strip() for ln in text.splitlines() if ln.strip()), "")
    m2 = re.match(r"^[A-Za-z0-9][A-Za-z0-9_.+-]*\s+\(([^)]+)\)\s*$", first)
    if m2:
        return (m2.group(1) or "").strip() or None
    return None


def check_package_update_available(package: str, timeout_seconds: float | int | None = 30) -> str | None:
    """Return latest PyPI version for an allowlisted MLX stack package if newer."""
    pkg = (package or "").strip()
    if pkg not in _MLX_TARGETED_UPGRADE_PACKAGES or not mlx_venv_exists():
        return None
    installed = get_installed_package_version(pkg)
    if not installed:
        return None
    try:
        r = subprocess.run(
            [_mlx_pip(), "index", "versions", pkg],
            capture_output=True, text=True, timeout=_pip_index_timeout_seconds(timeout_seconds, default=30),
        )
        latest = _pip_index_latest_version(r.stdout or "")
        if latest and latest != installed:
            return latest
    except Exception:
        pass
    return None


def get_mlx_stack_status(check_updates: bool = False) -> list[dict[str, str | None]]:
    """Return installed/update versions for the visible MLX package stack."""
    status: list[dict[str, str | None]] = []
    for pkg in _MLX_STACK_PACKAGE_NAMES:
        installed = get_installed_package_version(pkg)
        latest = check_package_update_available(pkg) if check_updates and installed else None
        status.append({
            "package": pkg,
            "installed": installed or "",
            "update_available": latest or "",
        })
    return status


def check_update_available(timeout_seconds: float | int | None = 30) -> str | None:
    """Return latest PyPI mlx-lm if newer than installed (for the 'Update to …' nudge only).

    This is **not** a claim that that mlx-lm version is safe with the current mlx-vlm, etc.:
    it only means PyPI has a newer mlx-lm. Users can still upgrade; use Revert to stable if
    the resolved stack misbehaves.
    """
    if not mlx_venv_exists():
        return None
    installed = get_installed_version()
    if not installed:
        return None
    try:
        r = subprocess.run(
            [_mlx_pip(), "index", "versions", "mlx-lm"],
            capture_output=True, text=True, timeout=_pip_index_timeout_seconds(timeout_seconds, default=30),
        )
        latest = _pip_index_latest_version(r.stdout or "")
        if latest and latest != installed:
            return latest
    except Exception:
        pass
    return None


def upgrade_mlx_stack_package(package: str) -> tuple[bool, str]:
    """Upgrade a single allowlisted MLX stack package in-place.

    This is intentionally allowlisted: the UI can expose targeted repairs such as
    mlx-openai-server without accepting arbitrary package names from a request body.
    """
    pkg = (package or "").strip()
    if pkg not in _MLX_TARGETED_UPGRADE_PACKAGES:
        return False, f"Package is not an allowlisted MLX stack package: {pkg}"
    if not mlx_venv_exists():
        return False, "MLX venv not installed."
    stop_mlx_server()
    try:
        subprocess.run(
            [_mlx_pip(), "install", "--upgrade", pkg],
            check=True, capture_output=True, text=True, timeout=_MLX_PIP_STACK_TIMEOUT_SEC,
        )
    except subprocess.CalledProcessError as e:
        return False, f"{pkg} upgrade failed: {e.stderr or e.stdout or str(e)}"
    except Exception as e:
        return False, f"{pkg} upgrade failed: {e}"
    _save_installed_version()
    mlx_prepare_venv_runtime()
    version = get_installed_package_version(pkg) or "unknown"
    return True, f"{pkg} upgraded to {version}. Load the selected model to restart MLX."


# ---------------------------------------------------------------------------
# Server lifecycle
# ---------------------------------------------------------------------------

_mlx_server_process: subprocess.Popen | None = None
_mlx_server_lock = Lock()


def mlx_server_running() -> bool:
    """Check if the local mlx-vlm server is reachable."""
    if not requests:
        return False
    try:
        r = requests.get(f"{MLX_SERVER_URL}/v1/models", timeout=2)
        return r.status_code == 200
    except Exception:
        return False


def _mlx_server_port_open() -> bool:
    """Return True when something is accepting TCP on the MLX port."""
    try:
        with socket.create_connection(("127.0.0.1", MLX_SERVER_PORT), timeout=1):
            return True
    except OSError:
        return False


def _tail_log(path: str, lines: int = 20) -> str:
    """Return the last N lines of a log file, or empty string if unreadable."""
    try:
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            all_lines = f.readlines()
        tail = "".join(all_lines[-lines:]).strip()
        return tail if tail else "(log is empty)"
    except Exception:
        return "(could not read log)"


def _resolve_mlx_model_path(model_path: str) -> str:
    """Return a model path that exists on disk for local paths.

    If model_path is an absolute path that does not exist, try the common
    typo Imstudio -> lmstudio (LM Studio uses .lmstudio and lmstudio-community).
    """
    path = (model_path or "").strip()
    if not path:
        return path
    if os.path.isdir(path):
        return path
    if "Imstudio" in path or "imstudio" in path:
        alt = path.replace("Imstudio", "lmstudio").replace("imstudio", "lmstudio")
        if os.path.isdir(alt):
            return alt
    return path


# Context length bounds for --context-length (match LLM prefs slider; reduces RAM when set lower).
_MLX_CONTEXT_LENGTH_MIN = 16384  # ~16k minimum across all models/engines
_MLX_CONTEXT_LENGTH_MAX = 1_048_576


# ---------------------------------------------------------------------------
# Native tool-call parser detection (advisory; mlx-vlm auto-infers at load time)
# ---------------------------------------------------------------------------

# Heuristic rules: (regex on lowercased model path/repo) → (tool_call_parser, reasoning_parser | None)
# Order matters — first match wins.  More specific patterns before generic ones.
_TOOL_CALL_PARSER_RULES: list[tuple[re.Pattern[str], str, str | None]] = [
    (re.compile(r"gemma[\-_]?4.*mlx"),    "gemma4",       None),
    (re.compile(r"gemma[\-_]?4"),         "gemma4",       None),
    (re.compile(r"qwen3[\.\-_]?6"),       "qwen3_coder",  "qwen3_5"),
    (re.compile(r"qwen3[\.\-_]?5"),       "qwen3_coder",  "qwen3_5"),
    (re.compile(r"qwen3[\.\-_]?coder"),   "qwen3_coder",  "qwen3"),
    (re.compile(r"qwen3"),                "qwen3",        "qwen3"),
    (re.compile(r"qwen2"),                "qwen3",        None),
    (re.compile(r"glm[\-_]?4"),           "glm4_moe",     None),
    (re.compile(r"minimax[\-_]?m2"),      "minimax_m2",   "minimax_m2"),
]

_KNOWN_TOOL_CALL_PARSERS = {
    # Names merged with ``_installed_parser_registry_names``; registry supplies model-specific parsers
    # (e.g. ``gemma4``). Typer enums vary by release — launch rejects unknown flags.
    "functiongemma",
    "glm4_moe",
    "harmony",
    "hermes",
    "iquest_coder_v1",
    "kimi_k2",
    "longcat_flash_lite",
    "minimax_m2",
    "nemotron3_nano",
    "qwen3",
    "qwen3_coder",
    "qwen3_moe",
    "qwen3_vl",
    "solar_open",
    "step_35",
}

_KNOWN_REASONING_PARSERS = {
    "glm4_moe",
    "glm47_flash",
    "hermes",
    "kimi_k2",
    "minimax_m2",
    "mixed_think_tool_handoff",
    "nemotron3_nano",
    "qwen3",
    "qwen3_5",
    "qwen3_moe",
    "qwen3_vl",
    "solar_open",
    "step_35",
}

_MAX_PARSER_OVERRIDES = 64
_MAX_OVERRIDE_KEY_CHARS = 512
_MAX_OVERRIDE_VALUE_CHARS = 64


def _installed_parser_registry_names(kind: str) -> set[str]:
    """Read available tool-parser names from the installed mlx-vlm / mlx-lm venv.

    mlx-vlm auto-infers the tool parser from a model's chat template and loads parser
    modules dynamically from ``mlx_vlm.tool_parsers`` / ``mlx_lm.tool_parsers``; there is
    no launch-time parser flag anymore. This introspection is now advisory/diagnostic
    only (it feeds parser-availability hints), so enumerate the parser submodules of those
    two packages. mlx has no separate reasoning-parser registry, so ``reasoning`` returns
    an empty set (callers still merge the static ``_KNOWN_REASONING_PARSERS``).
    """
    if kind != "tool":
        return set()
    if not mlx_venv_exists():
        return set()
    py = _mlx_python()
    if not os.path.isfile(py):
        return set()
    code = (
        "import json, pkgutil, importlib\n"
        "names = set()\n"
        "for mod in ('mlx_vlm.tool_parsers', 'mlx_lm.tool_parsers'):\n"
        "    try:\n"
        "        pkg = importlib.import_module(mod)\n"
        "        for m in pkgutil.iter_modules(pkg.__path__):\n"
        "            if not m.name.startswith('_'):\n"
        "                names.add(m.name)\n"
        "    except Exception:\n"
        "        pass\n"
        "print(json.dumps(sorted(names)))\n"
    )
    try:
        r = subprocess.run(
            [py, "-c", code],
            capture_output=True, text=True, timeout=10,
        )
        if r.returncode != 0:
            return set()
        data = json.loads((r.stdout or "").strip() or "[]")
        if isinstance(data, list):
            return {str(name).strip().lower() for name in data if str(name).strip()}
    except Exception:
        logger.debug("Failed to inspect mlx-vlm/mlx-lm parser registry", exc_info=True)
    return set()


def _available_tool_call_parsers() -> set[str]:
    return set(_KNOWN_TOOL_CALL_PARSERS) | _installed_parser_registry_names("tool")


def _available_reasoning_parsers() -> set[str]:
    return set(_KNOWN_REASONING_PARSERS) | _installed_parser_registry_names("reasoning")


def _compatible_parser_names(tcp: str | None, rp: str | None) -> tuple[str | None, str | None, str]:
    """Validate parser names against the installed server, with Gemma 4 fallback."""
    tcp_norm = _normalize_parser_name(tcp)
    rp_norm = _normalize_parser_name(rp)
    tool_parsers = _available_tool_call_parsers()
    reasoning_parsers = _available_reasoning_parsers()

    warning = ""
    if tcp_norm and tcp_norm not in tool_parsers:
        if tcp_norm == "gemma4" and "functiongemma" in tool_parsers:
            tcp_norm = "functiongemma"
            warning = "gemma4 parser is unavailable in this mlx-openai-server; falling back to functiongemma"
        else:
            return None, rp_norm or None, f"tool_call_parser unavailable in installed mlx-openai-server: {tcp_norm}"

    if rp_norm and rp_norm not in reasoning_parsers:
        return tcp_norm or None, None, f"reasoning_parser unavailable in installed mlx-openai-server: {rp_norm}"

    # When ``gemma4`` is advertised by the installed server, prefer it — it matches MLX_RECOMMENDED for Gemma 4 and
    # aligns with mlx-openai-server parsers that emit native streaming ``delta.tool_calls``. Older builds that
    # lack ``gemma4`` fall through above to ``functiongemma``.

    return tcp_norm or None, rp_norm or None, warning


@dataclass(frozen=True)
class MlxParserResolution:
    tool_call_parser: str | None
    reasoning_parser: str | None
    source: str
    reason: str
    normalized_model: str
    matched_key: str = ""
    warning: str = ""

    def to_dict(self) -> dict[str, str | None]:
        return asdict(self)


def _load_app_prefs_for_mlx() -> dict:
    """Load prefs from ``APP_PREFS_PATH`` so tests can monkeypatch the path.

    ``load_app_prefs_cached`` always reads the canonical store and ignores
    ``app_mlx.APP_PREFS_PATH``, which breaks unit tests that write isolated JSON.
    """
    try:
        path = APP_PREFS_PATH
        if path and os.path.isfile(path):
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, dict):
                return data
    except Exception:
        logger.debug("Failed to read app prefs for MLX parser overrides", exc_info=True)
    return {}


def _normalize_mlx_model_identity(model_path: str) -> str:
    path = _resolve_mlx_model_path(model_path or "").strip()
    if not path:
        return ""
    expanded = os.path.expanduser(path)
    if os.path.isabs(expanded):
        return os.path.realpath(expanded).rstrip("/")
    return path.strip().rstrip("/").lower()


_HF_CACHE_DIR_TOKEN_RE = re.compile(r"^models--(?P<org>.+)--(?P<name>.+)$")
_HF_SNAPSHOT_HASH_RE = re.compile(r"^[a-f0-9]{40}$")


def _mlx_model_slug(model_path: str) -> str:
    """Filesystem-safe slug for adapter names (model id, not snapshot revision hash)."""
    ident = _normalize_mlx_model_identity(model_path)
    if not ident:
        return ""
    parts = [p for p in ident.split("/") if p]
    for part in reversed(parts):
        m = _HF_CACHE_DIR_TOKEN_RE.match(part)
        if m:
            return m.group("name").lower()
        if "__" in part:
            org_name = part.split("__", 1)
            if len(org_name) == 2 and org_name[1]:
                return org_name[1].lower()
    if len(parts) >= 2 and _HF_SNAPSHOT_HASH_RE.match(parts[-1]):
        parent = parts[-2]
        m = _HF_CACHE_DIR_TOKEN_RE.match(parent)
        if m:
            return m.group("name").lower()
        if "__" in parent:
            org_name = parent.split("__", 1)
            if len(org_name) == 2 and org_name[1]:
                return org_name[1].lower()
    last = parts[-1].lower() if parts else ""
    if len(parts) == 2 and not ident.startswith("/"):
        return last
    return last


def _normalize_parser_name(value: object) -> str:
    return str(value or "").strip().lower()[:_MAX_OVERRIDE_VALUE_CHARS]


def _override_parser_names(raw_value: object) -> tuple[str | None, str | None, str]:
    if isinstance(raw_value, str):
        tcp = _normalize_parser_name(raw_value)
        return (tcp or None), None, ""
    if not isinstance(raw_value, dict):
        return None, None, "override value must be a string or object"
    tcp = _normalize_parser_name(raw_value.get("tool_call_parser") or raw_value.get("parser"))
    rp = _normalize_parser_name(raw_value.get("reasoning_parser"))
    return (tcp or None), (rp or None), ""


def _validated_override(raw_value: object) -> tuple[str | None, str | None, str]:
    tcp, rp, error = _override_parser_names(raw_value)
    if error:
        return None, None, error
    if not tcp:
        return None, None, "override is missing tool_call_parser"
    tcp, rp, compat_note = _compatible_parser_names(tcp, rp)
    if tcp is None:
        # Hard failure (unknown parser): ``compat_note`` is the rejection reason.
        return None, None, compat_note or "invalid tool_call_parser"
    # Successful resolution, including Gemma 4 downgrade ``gemma4`` → ``functiongemma`` — that path sets a
    # non‑empty informational ``compat_note`` which must not invalidate the override.
    return tcp, rp, ""


def _parser_override_candidates(model_path: str) -> list[str]:
    ident = _normalize_mlx_model_identity(model_path)
    slug = _mlx_model_slug(model_path)
    candidates = []
    if ident:
        candidates.append(ident)
        candidates.append(ident.lower())
    if slug and slug not in candidates:
        candidates.append(slug)
    return candidates


def _parser_override_for_model(model_path: str) -> MlxParserResolution | None:
    prefs = _load_app_prefs_for_mlx()
    raw_overrides = prefs.get("mlx_tool_call_parser_overrides")
    if not isinstance(raw_overrides, dict) or not raw_overrides:
        return None
    if len(raw_overrides) > _MAX_PARSER_OVERRIDES:
        return MlxParserResolution(
            None,
            None,
            "override",
            "ignored parser overrides because there are too many entries",
            _normalize_mlx_model_identity(model_path),
            warning="too_many_overrides",
        )

    normalized_overrides: dict[str, object] = {}
    for key, value in raw_overrides.items():
        key_s = str(key or "").strip()
        if not key_s or len(key_s) > _MAX_OVERRIDE_KEY_CHARS:
            continue
        norm_key = _normalize_mlx_model_identity(key_s)
        if norm_key:
            normalized_overrides[norm_key] = value
            normalized_overrides[norm_key.lower()] = value
        slug_key = norm_key.rsplit("/", 1)[-1].lower() if "/" in norm_key else norm_key.lower()
        if slug_key:
            normalized_overrides.setdefault(slug_key, value)

    for candidate in _parser_override_candidates(model_path):
        if candidate not in normalized_overrides:
            continue
        tcp, rp, error = _validated_override(normalized_overrides[candidate])
        if error:
            return MlxParserResolution(
                None,
                None,
                "override",
                f"ignored parser override for {candidate}: {error}",
                _normalize_mlx_model_identity(model_path),
                matched_key=candidate,
                warning=error,
            )
        return MlxParserResolution(
            tcp,
            rp,
            "override",
            f"matched user parser override for {candidate}",
            _normalize_mlx_model_identity(model_path),
            matched_key=candidate,
        )
    return None


def resolve_mlx_tool_call_parsers(model_path: str) -> MlxParserResolution:
    """Resolve native tool parser metadata for a model path/repo.

    User overrides win, followed by curated recommended models, then heuristics.
    Unknown models intentionally return no parser so XML fallback can still work.
    """
    normalized_model = _normalize_mlx_model_identity(model_path)
    override = _parser_override_for_model(model_path)
    if override and (override.tool_call_parser or override.warning):
        return override

    for entry in MLX_RECOMMENDED:
        repo = str(entry.get("repo", "") or "").strip()
        if not repo:
            continue
        if not _catalog_repo_matches_model_path(repo, model_path):
            continue
        tcp = entry.get("tool_call_parser")
        rp = entry.get("reasoning_parser")
        if tcp:
            resolved_tcp, resolved_rp, warning = _compatible_parser_names(str(tcp), str(rp) if rp else None)
            reason = f"matched recommended model {repo}"
            if warning:
                reason = f"{reason}; {warning}"
            return MlxParserResolution(
                resolved_tcp,
                resolved_rp,
                "recommended",
                reason,
                normalized_model,
                matched_key=repo,
                warning=warning,
            )
        break

    slug = _mlx_model_slug(model_path)
    for pattern, tcp, rp in _TOOL_CALL_PARSER_RULES:
        if pattern.search(slug):
            resolved_tcp, resolved_rp, warning = _compatible_parser_names(tcp, rp)
            reason = f"matched model slug heuristic {pattern.pattern}"
            if warning:
                reason = f"{reason}; {warning}"
            return MlxParserResolution(
                resolved_tcp,
                resolved_rp,
                "heuristic",
                reason,
                normalized_model,
                matched_key=pattern.pattern,
                warning=warning,
            )

    return MlxParserResolution(
        None,
        None,
        "none",
        "no known native tool-call parser for this model",
        normalized_model,
    )


def _detect_tool_call_parsers(model_path: str) -> tuple[str | None, str | None]:
    """Return (tool_call_parser, reasoning_parser) for *model_path*.

    1. Exact match against MLX_RECOMMENDED (by repo field).
    2. Heuristic regex match on the lowercased model path.
    3. (None, None) — server runs without native tool-call parsing; Mugi's
       XML/channel salvage will still work as a fallback.
    """
    resolved = resolve_mlx_tool_call_parsers(model_path)
    return resolved.tool_call_parser, resolved.reasoning_parser


def _catalog_repo_matches_model_path(repo: str, model_path: str) -> bool:
    """True when *model_path* refers to catalog repo (HF id, Mugi ``org__repo``, or HF cache)."""
    repo = str(repo or "").strip()
    if not repo or not str(model_path or "").strip():
        return False
    repo_lower = repo.lower()
    repo_slug = repo_lower.split("/", 1)[-1]
    org, name = repo_lower.split("/", 1) if "/" in repo_lower else ("", repo_lower)
    mugi_token = f"{org}__{name}" if org else name.replace("/", "__")
    hf_token = f"models--{org}--{name}" if org else ""

    haystacks = {
        (model_path or "").lower().rstrip("/"),
        _normalize_mlx_model_identity(model_path).lower(),
        _mlx_model_slug(model_path).lower(),
    }
    for hay in haystacks:
        if not hay:
            continue
        if hay == repo_lower or hay.endswith("/" + repo_slug) or hay.endswith("/" + mugi_token):
            return True
        tail = hay.rsplit("/", 1)[-1]
        if tail == mugi_token or tail.replace("__", "/") == repo_lower:
            return True
        if hf_token and hf_token in hay:
            return True
    return False


def _catalog_entry_for_model_path(model_path: str) -> dict | None:
    """Return the MLX_RECOMMENDED row whose repo matches *model_path* (HF id or local path)."""
    for entry in MLX_RECOMMENDED:
        repo = str(entry.get("repo", "") or "").strip()
        if repo and _catalog_repo_matches_model_path(repo, model_path):
            return entry
    return None


def resolve_recommended_model_type(model_path: str, model_type: str = "auto") -> str:
    """Resolve the chat model type (``text`` / ``multimodal``) for the fingerprint/UI.

    mlx-vlm auto-detects model type at load, so this is no longer a launch flag — it is
    retained for the runtime fingerprint and Settings display. Explicit ``text`` /
    ``multimodal`` are honored. For ``auto``, catalog entries with ``vision`` in caps route
    multimodal even before ``config.json`` exists on disk (setup wizard passes HF repo ids
    pre-download).
    """
    mt = (model_type or "auto").strip().lower() or "auto"
    if mt in ("text", "multimodal"):
        return mt
    resolved = _resolve_mlx_model_path(model_path)
    entry = _catalog_entry_for_model_path(resolved or model_path)
    if entry and "vision" in (entry.get("caps") or []):
        return "multimodal"
    return _detect_model_type(resolved or model_path)


def effective_mlx_model_type(model_path: str, model_type: str = "auto") -> str:
    return resolve_recommended_model_type(model_path, model_type)


def build_mlx_runtime_fingerprint(
    model_path: str,
    model_type: str = "auto",
    context_length: int | None = None,
) -> dict:
    resolved_model_path = _resolve_mlx_model_path(model_path)
    effective_type = effective_mlx_model_type(resolved_model_path, model_type)
    parser_resolution = resolve_mlx_tool_call_parsers(resolved_model_path)
    ctx = None
    try:
        if context_length is not None:
            ctx_int = int(context_length)
            if _MLX_CONTEXT_LENGTH_MIN <= ctx_int <= _MLX_CONTEXT_LENGTH_MAX:
                ctx = ctx_int
    except (TypeError, ValueError):
        ctx = None
    return {
        "model_path": resolved_model_path,
        "normalized_model": parser_resolution.normalized_model,
        "model_type": effective_type,
        "context_length": ctx,
        "tool_call_parser": parser_resolution.tool_call_parser,
        "reasoning_parser": parser_resolution.reasoning_parser,
        "parser": parser_resolution.to_dict(),
        "mlx_lm_version": get_installed_version() or "",
    }


def _runtime_fingerprint_core(fp: dict | None) -> dict:
    # Restart decision keys on the actual mlx-vlm launch inputs only: the preloaded model
    # and --max-kv-size (context_length). model_type and tool/reasoning parsers are no longer
    # launch flags (mlx-vlm auto-detects them), so they stay in the full fingerprint for
    # UI/diagnostics but must NOT force a teardown of a healthy server when they change.
    fp = fp or {}
    return {
        "normalized_model": fp.get("normalized_model") or "",
        "context_length": fp.get("context_length"),
    }


def runtime_fingerprint_matches(expected: dict | None, actual: dict | None) -> bool:
    return _runtime_fingerprint_core(expected) == _runtime_fingerprint_core(actual)


def read_mlx_runtime_state() -> dict:
    try:
        if os.path.isfile(MLX_RUNTIME_STATE_FILE):
            with open(MLX_RUNTIME_STATE_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, dict):
                return data
    except Exception:
        logger.debug("Failed to read MLX runtime state", exc_info=True)
    return {}


def get_last_mlx_runtime_fingerprint() -> dict:
    state = read_mlx_runtime_state()
    fp = state.get("fingerprint") if isinstance(state, dict) else None
    return fp if isinstance(fp, dict) else {}


def _write_mlx_runtime_state(fingerprint: dict, status: str, message: str = "") -> None:
    os.makedirs(MLX_DATA_DIR, exist_ok=True)
    payload = {
        "status": status,
        "message": message,
        "timestamp": int(time.time()),
        "fingerprint": fingerprint,
    }
    tmp = MLX_RUNTIME_STATE_FILE + ".tmp"
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2, sort_keys=True)
            f.write("\n")
        os.chmod(tmp, 0o600)
        os.replace(tmp, MLX_RUNTIME_STATE_FILE)
    except Exception:
        logger.debug("Failed to write MLX runtime state", exc_info=True)
        try:
            if os.path.exists(tmp):
                os.unlink(tmp)
        except Exception:
            pass


def mlx_adapter_path(adapter_name: str) -> str:
    """Absolute path to a stored LoRA adapter directory.

    mlx-vlm reads a per-request ``adapter_path`` field; adapters live under MLX_ADAPTERS_DIR.
    """
    return os.path.join(MLX_ADAPTERS_DIR, (adapter_name or "").strip())


_MLX_LAUNCH_EXTRA_BLOCKED_FLAGS = {"--host", "--trust-remote-code"}


def _sanitize_mlx_launch_extra(tokens: list[str]) -> list[str]:
    """Drop security-sensitive overrides from MUGI_MLX_LAUNCH_EXTRA.

    MUGI_MLX_LAUNCH_EXTRA is appended AFTER our own flags, so without this an
    operator value like ``--host 0.0.0.0`` or ``--trust-remote-code`` would
    silently override the loopback-only / no-remote-code security defaults.
    Strips both flag forms (``--host x`` and ``--host=x``) and the
    ``--trust-remote-code`` store_true flag.
    """
    out: list[str] = []
    skip_next = False
    for tok in tokens:
        if skip_next:
            skip_next = False
            continue
        base = tok.split("=", 1)[0]
        if base in _MLX_LAUNCH_EXTRA_BLOCKED_FLAGS:
            # `--host value` (space form) consumes the following token; `--host=value`
            # and store_true flags consume nothing extra.
            if base == "--host" and "=" not in tok:
                skip_next = True
            logger.warning("Ignoring blocked MUGI_MLX_LAUNCH_EXTRA flag: %s", base)
            continue
        out.append(tok)
    return out


def start_mlx_server(
    model_path: str,
    model_type: str = "auto",
    context_length: int | None = None,
) -> tuple[bool, str]:
    """Start the local mlx-vlm server (``python -m mlx_vlm.server``) as a subprocess.

    model_path: HuggingFace repo ID or local path to the model directory.
    model_type: "auto"/"text"/"multimodal" — retained for the fingerprint/UI only;
        mlx-vlm auto-detects model type and tool parser, so it is no longer a launch flag.
    context_length: the user's context-window preference. It drives CLIENT-SIDE context
        trimming (agent_loop_trim_context), NOT a server launch flag. We deliberately do NOT
        map it to --max-kv-size: mlx-vlm has no --context-length, and --max-kv-size is a
        *rotating KV cache cap* that (a) raises PromptTooLongError when prompt+max_tokens
        exceed it (so a 32k cap hard-fails Mugi's own system prompt, which needs >=64k) and
        (b) silently drops the oldest tokens once exceeded. The model config drives the real
        window; power users who genuinely want a memory cap can pass
        MUGI_MLX_LAUNCH_EXTRA='--max-kv-size N'.
    Returns (success, message).
    """
    global _mlx_server_process
    with _mlx_server_lock:
        if not mlx_venv_exists():
            return False, "MLX venv not installed."

        py = _mlx_python()
        # bin/mlx_vlm.server is the 0.6.x console script; we launch via `-m` for robustness
        # but use its presence as a legible "stack is current" signal.
        server_script = os.path.join(MLX_VENV_DIR, "bin", "mlx_vlm.server")
        if not os.path.isfile(py) or not os.path.isfile(server_script):
            return False, (
                "mlx-vlm server not found in the MLX venv. Reinstall the MLX stack "
                "from Settings → LLMs (the local server is now python -m mlx_vlm.server)."
            )

        # pip upgrades replace site-packages; git pull does not. Re-apply vendored patches on
        # every server start so multimodal MLX (asyncio-thread mx.array) keeps working without
        # a manual Settings reinstall.
        patched_site = mlx_prepare_venv_runtime()

        # Resolve local path so it exists (e.g. .Imstudio -> .lmstudio)
        model_path = _resolve_mlx_model_path(model_path)
        # If path looks like a local absolute path, it must exist or the server will fail with "Repo id must be..."
        if model_path.startswith("/") or model_path.startswith(os.path.expanduser("~")):
            if not os.path.isdir(model_path):
                return False, (
                    f"Model path does not exist or is not a directory: {model_path}. "
                    "If this is an LM Studio model, check that the path uses .lmstudio (lowercase L), not .Imstudio."
                )

        fingerprint = build_mlx_runtime_fingerprint(model_path, model_type, context_length)
        model_type = str(fingerprint.get("model_type") or "text")
        # Preload under the SAME normalized identity the agent sends as the request `model`
        # (get_mlx_openai_model_id -> _normalize_mlx_model_identity), so requests hit the
        # preloaded cache key with zero reloads.
        launch_model = _normalize_mlx_model_identity(model_path) or model_path

        if mlx_server_running():
            current_fp = get_last_mlx_runtime_fingerprint()
            _fp_match = runtime_fingerprint_matches(fingerprint, current_fp)
            logger.debug("MLX start fp check: match=%s patched=%s model=%s",
                         _fp_match, patched_site, str(fingerprint.get("model_path", ""))[-80:])
            if _fp_match and not patched_site:
                return True, "MLX server already running with matching model."
            _stop_mlx_server_unlocked()

        # mlx-vlm auto-detects model type via config and auto-infers the tool-call parser
        # from the chat template, so --model-type / --tool-call-parser / --reasoning-parser
        # are gone. --host 127.0.0.1 is load-bearing (mlx-vlm defaults to 0.0.0.0 = LAN-exposed);
        # we never pass --trust-remote-code.
        cmd = [
            py, "-m", "mlx_vlm.server",
            "--host", "127.0.0.1",
            "--port", str(MLX_SERVER_PORT),
            "--model", launch_model,
        ]
        # NOTE: we intentionally do NOT pass --max-kv-size derived from the context-window
        # preference. mlx-vlm treats --max-kv-size as a hard budget: prompt+max_tokens over it
        # raises PromptTooLongError (a 32k cap broke chat because Mugi's own context needs
        # >=64k) and it builds a truncating rotating cache. The model config drives the window;
        # context_length is enforced client-side via trimming. A genuine memory cap can still be
        # set with MUGI_MLX_LAUNCH_EXTRA='--max-kv-size N' (see _sanitize_mlx_launch_extra).

        # Parser/model_type are now metadata (logged below for diagnostics), not launch flags.
        _parser = fingerprint.get("parser") if isinstance(fingerprint.get("parser"), dict) else {}
        _tcp = fingerprint.get("tool_call_parser")
        _rp = fingerprint.get("reasoning_parser")

        # Extra `python -m mlx_vlm.server` args (e.g. experimental upstream flags).
        # Example: MUGI_MLX_LAUNCH_EXTRA='--prefill-step-size 512'  (shlex-split; be careful
        # with quotes). --host / --trust-remote-code are stripped (security defaults win).
        _extra = (os.environ.get("MUGI_MLX_LAUNCH_EXTRA") or "").strip()
        if _extra:
            cmd.extend(_sanitize_mlx_launch_extra(shlex.split(_extra)))

        # Automatic Prompt Caching (APC). mlx-vlm's server defaults APC OFF
        # (apc.from_env: APC_ENABLED defaults to "0"), so WITHOUT this every chat turn
        # re-prefills the ENTIRE prompt (system prompt + the "You"/User Model block + memory +
        # tools + history) from scratch — measured ~7s at 20k tokens and ~20s at 40k tokens of
        # time-to-first-token, paid on EVERY turn with no follow-up reuse. The retired
        # mlx-openai-server cached the KV prefix so follow-ups were near-instant; enabling APC
        # restores that. APC reuse is *exact* (byte-identical K/V keyed on the token prefix), and
        # Mugi binds a single matching LoRA adapter per loaded model for the whole server lifetime
        # (the adapter<->model gate + probe "inherit"), so cross-adapter KV reuse cannot occur.
        # Tunable: set APC_ENABLED=0 in the environment to opt out; pin the pool with
        # MUGI_MLX_APC_NUM_BLOCKS (or APC_NUM_BLOCKS). Each block is 16 tokens. The pool is
        # PRE-ALLOCATED up front, so a fixed 64k-token pool (4096 blocks) that is free on a
        # 128 GB box is several GB of dead-weight KV on a 24 GB box — enough to tip generation
        # into swap. When the size isn't pinned we scale it to unified RAM (~16k tokens on
        # <=16 GB up to 64k on >48 GB); a pool smaller than the prompt just caches less of the
        # prefix, it never breaks correctness. See _apc_num_blocks_for_ram().
        server_env = dict(os.environ)
        server_env.setdefault("APC_ENABLED", "1")
        if server_env.get("APC_ENABLED", "0") in ("1", "true", "True", "yes") and "APC_NUM_BLOCKS" not in server_env:
            _apc_blocks = (os.environ.get("MUGI_MLX_APC_NUM_BLOCKS") or "").strip()
            server_env["APC_NUM_BLOCKS"] = _apc_blocks if _apc_blocks.isdigit() else str(_apc_num_blocks_for_ram())

        os.makedirs(MLX_DATA_DIR, exist_ok=True)
        try:
            with open(os.path.join(MLX_DATA_DIR, "startup_resolution.log"), "a", encoding="utf-8") as f:
                f.write(
                    f"{int(time.time())} model={launch_model} resolved_model_type={model_type} "
                    f"context_length={fingerprint.get('context_length') or ''} "
                    f"tool_call_parser={_tcp or ''} reasoning_parser={_rp or ''} "
                    f"parser_source={_parser.get('source') or ''} parser_reason={_parser.get('reason') or ''} "
                    f"apc_enabled={server_env.get('APC_ENABLED')} apc_num_blocks={server_env.get('APC_NUM_BLOCKS', '')}\n"
                )
        except Exception:
            pass
        log_path = os.path.join(MLX_DATA_DIR, "server.log")
        try:
            log_file = open(log_path, "w")
            _mlx_server_process = subprocess.Popen(
                cmd,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                cwd=MLX_DATA_DIR,
                env=server_env,
            )
            _write_mlx_runtime_state(fingerprint, "starting", "MLX server process launched.")
        except Exception as e:
            return False, f"Failed to start server: {e}"

        # MLX model loading can be slow on first launch (shader compilation + weight loading).
        # mlx-vlm preloads the model in a BLOCKING FastAPI lifespan, so the HTTP port only
        # opens AFTER the weights finish loading — there is no early "port open but warming"
        # window like mlx-openai-server had. So we treat "process still alive" (poll() is None)
        # as warming and never hard-fail a healthy-but-slow load.
        for _ in range(90):
            time.sleep(1)
            if mlx_server_running():
                _write_mlx_runtime_state(fingerprint, "running", "MLX server started.")
                return True, "MLX server started."
            if _mlx_server_process.poll() is not None:
                return False, f"MLX server exited unexpectedly.\n{_tail_log(log_path)}"
        # 90s elapsed without a 200 from /v1/models. If the process is still alive it is almost
        # certainly still loading a large model (lifespan preload), so report warming rather than
        # a false failure; chat verification retries will confirm readiness shortly after.
        if _mlx_server_process and _mlx_server_process.poll() is None:
            _write_mlx_runtime_state(fingerprint, "warming", "MLX server is loading the model in the background.")
            return True, "MLX server is running and finishing model load in the background."
        _write_mlx_runtime_state(fingerprint, "failed", "MLX server did not become ready within 90 seconds.")
        return False, f"MLX server did not become ready within 90 seconds.\n{_tail_log(log_path)}"


def _stop_mlx_server_unlocked():
    """Stop the mlx-vlm server subprocess. Caller must hold _mlx_server_lock."""
    logger.debug("Stopping MLX server (unlocked)")
    global _mlx_server_process
    proc = _mlx_server_process
    if proc is not None:
        try:
            proc.terminate()
            proc.wait(timeout=10)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass
        _mlx_server_process = None
    # Also try to kill any orphan on the port
    _kill_process_on_port(MLX_SERVER_PORT)


def stop_mlx_server():
    """Stop the mlx-vlm server subprocess."""
    with _mlx_server_lock:
        _stop_mlx_server_unlocked()
        state = read_mlx_runtime_state()
        fp = state.get("fingerprint") if isinstance(state, dict) else {}
        if isinstance(fp, dict):
            _write_mlx_runtime_state(fp, "stopped", "MLX server stopped.")


def _kill_process_on_port(port: int):
    """Best-effort kill of any process listening on the given port."""
    try:
        r = subprocess.run(
            ["lsof", "-ti", f":{port}"],
            capture_output=True, text=True, timeout=5,
        )
        for pid_str in (r.stdout or "").split():
            try:
                os.kill(int(pid_str), signal.SIGTERM)
            except (ValueError, OSError):
                pass
    except Exception:
        pass


def reinstall_mlx_venv_and_start(
    model_path: str = "",
    context_length: int | None = None,
) -> tuple[bool, str]:
    """Stop the server, remove the venv (keeping models/adapters), reinstall, then start.

    Returns (success, message). If model_path is empty the server is not started after install.
    """
    stop_mlx_server()
    ok, msg = delete_mlx_venv(remove_data=False)
    if not ok:
        return False, f"Failed to remove old MLX venv: {msg}"
    ok, msg = ensure_mlx_venv(prefer_pinned=True)
    if not ok:
        return False, f"Failed to reinstall MLX: {msg}"
    new_version = get_installed_version() or "unknown"
    pinned = get_pinned_mlx_stack_summary()
    stack_note = f" ({pinned})" if pinned else ""
    if not model_path:
        return True, (
            f"MLX reinstalled to pinned stable stack{stack_note}, mlx-lm {new_version}. "
            "No model configured; server not started."
        )
    ok_health, health_msg = _restart_selected_model_and_verify_chat(
        f"MLX reinstalled (mlx-lm {new_version})."
    )
    if ok_health:
        return True, health_msg
    if "could not restart" in health_msg:
        return False, health_msg
    return True, _revert_verify_soft_success_message(new_version, health_msg)


# ---------------------------------------------------------------------------
# Model discovery
# ---------------------------------------------------------------------------


def _is_mlx_model_dir(path: str) -> bool:
    """Return True if path looks like an MLX-format model (has safetensors, no GGUF)."""
    if not os.path.isdir(path):
        return False
    has_safetensors = False
    has_gguf = False
    try:
        for entry in os.scandir(path):
            name_lower = entry.name.lower()
            if name_lower.endswith(".safetensors"):
                has_safetensors = True
            if name_lower.endswith(".gguf"):
                has_gguf = True
            if has_safetensors:
                break
    except OSError:
        return False
    return has_safetensors and not has_gguf


def _detect_model_type(model_path: str) -> str:
    """Inspect config.json to determine if a model is multimodal (vision) or text-only."""
    config_path = os.path.join(model_path, "config.json") if os.path.isdir(model_path) else ""
    if not config_path or not os.path.isfile(config_path):
        return "text"
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = json.load(f)
    except Exception:
        return "text"
    model_type = (cfg.get("model_type") or "").lower()
    architectures = [a.lower() for a in (cfg.get("architectures") or [])]
    vision_keywords = ("vl", "vision", "visual", "multimodal", "image")
    if any(kw in model_type for kw in vision_keywords):
        return "multimodal"
    if any(any(kw in arch for kw in vision_keywords) for arch in architectures):
        return "multimodal"
    if cfg.get("vision_config") or cfg.get("visual"):
        return "multimodal"
    return "text"


def mlx_find_compatible_models() -> list[dict]:
    """Scan LM Studio model dirs and HuggingFace cache for MLX-format models.
    Returns list of {"path": str, "display_name": str, "source": "lmstudio"|"huggingface"}.
    """
    found: list[dict] = []
    seen_paths: set[str] = set()

    for base_dir in _LMSTUDIO_MODEL_DIRS:
        if not os.path.isdir(base_dir):
            continue
        _scan_model_tree(base_dir, found, seen_paths, source="lmstudio")

    hf_cache = os.path.expanduser("~/.cache/huggingface/hub")
    if os.path.isdir(hf_cache):
        _scan_hf_cache(hf_cache, found, seen_paths)
    _scan_mugi_models_dir(found, seen_paths)

    found.sort(key=lambda m: m["display_name"].lower())
    return found


def _scan_model_tree(base_dir: str, out: list, seen: set, source: str):
    """Walk two levels deep looking for MLX model directories."""
    try:
        for publisher in os.scandir(base_dir):
            if not publisher.is_dir():
                continue
            for model in os.scandir(publisher.path):
                if not model.is_dir():
                    continue
                real = os.path.realpath(model.path)
                if real in seen:
                    continue
                if _is_mlx_model_dir(model.path):
                    seen.add(real)
                    display = f"{publisher.name}/{model.name}"
                    out.append({"path": model.path, "display_name": display, "source": source})
    except OSError:
        pass


def _scan_hf_cache(hf_cache: str, out: list, seen: set):
    """Scan HuggingFace cache for MLX-format models (models--org--name/snapshots/*)."""
    try:
        for entry in os.scandir(hf_cache):
            if not entry.name.startswith("models--") or not entry.is_dir():
                continue
            parts = entry.name.split("--", 2)
            if len(parts) < 3:
                continue
            org, name = parts[1], parts[2]
            snapshots_dir = os.path.join(entry.path, "snapshots")
            if not os.path.isdir(snapshots_dir):
                continue
            for snap in os.scandir(snapshots_dir):
                if not snap.is_dir():
                    continue
                real = os.path.realpath(snap.path)
                if real in seen:
                    continue
                if _is_mlx_model_dir(snap.path):
                    seen.add(real)
                    out.append({
                        "path": snap.path,
                        "display_name": f"{org}/{name}",
                        "source": "huggingface",
                    })
                    break  # one snapshot per model is enough
    except OSError:
        pass


def _scan_mugi_models_dir(out: list, seen: set):
    """Scan canonical ~/.mugi/mlx/models paths created after download finalization."""
    if not os.path.isdir(MLX_CANONICAL_MODELS_DIR):
        return
    try:
        for entry in os.scandir(MLX_CANONICAL_MODELS_DIR):
            if not (entry.is_symlink() or entry.is_dir(follow_symlinks=True)):
                continue
            real = os.path.realpath(entry.path)
            if real in seen:
                continue
            if _is_mlx_model_dir(entry.path):
                seen.add(real)
                out.append({
                    "path": entry.path,
                    "display_name": entry.name.replace("__", "/"),
                    "source": "mugi",
                })
    except OSError:
        pass


def get_mlx_recommended_models() -> list[dict]:
    """Return recommended MLX models list.

    Guarantees:
    - Only models with tool-calling and vision support are listed (required for Mugi's agent features).
    - RAM-tier recommendation is surfaced via `recommended: bool`.
    - The current-machine recommended model is first in the list.
    """
    models = [
        dict(m)
        for m in MLX_RECOMMENDED
        if "tools" in (m.get("caps") or []) and "vision" in (m.get("caps") or [])
    ]
    ram_gb = _get_system_ram_gb()

    for m in models:
        min_gb = m.get("recommended_for_min_ram_gb")
        max_gb = m.get("recommended_for_max_ram_gb")
        is_recommended = False
        if isinstance(ram_gb, int):
            if isinstance(min_gb, int) and isinstance(max_gb, int):
                is_recommended = min_gb <= ram_gb <= max_gb
            elif isinstance(min_gb, int):
                is_recommended = ram_gb >= min_gb
            elif isinstance(max_gb, int):
                is_recommended = ram_gb <= max_gb
        m["recommended"] = is_recommended

    # Place recommended models first while preserving relative order.
    recommended = [m for m in models if m.get("recommended") is True]
    others = [m for m in models if m.get("recommended") is not True]
    return recommended + others


def hf_model_repo_exists(repo_id: str, *, timeout: float = 20.0) -> bool:
    """Return True when *repo_id* resolves on Hugging Face (download target is valid)."""
    repo_id = str(repo_id or "").strip()
    if not repo_id or "/" not in repo_id:
        return False
    try:
        import urllib.request

        url = f"https://huggingface.co/api/models/{repo_id}"
        req = urllib.request.Request(url, headers={"User-Agent": "Mugi-ModelVerify/1"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return 200 <= int(getattr(resp, "status", 200) or 200) < 300
    except Exception:
        return False


def mlx_recommended_hf_repos() -> list[str]:
    """HF repo ids from ``MLX_RECOMMENDED`` (deduped, stable order)."""
    seen: set[str] = set()
    out: list[str] = []
    for entry in MLX_RECOMMENDED:
        repo = str(entry.get("repo") or "").strip()
        if repo and repo not in seen:
            seen.add(repo)
            out.append(repo)
    return out


def _get_system_ram_gb() -> int | None:
    """Return physical RAM in GiB when available."""
    try:
        if sys.platform == "darwin":
            out = subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True, timeout=2).strip()
            if out.isdigit():
                return int(int(out) / (1024 ** 3))
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
        if isinstance(pages, int) and isinstance(page_size, int) and pages > 0 and page_size > 0:
            return int((pages * page_size) / (1024 ** 3))
    except Exception:
        return None
    return None


def _apc_num_blocks_for_ram() -> int:
    """Pick the APC KV-pool size (in 16-token blocks) for the current machine.

    mlx-vlm pre-allocates the APC pool up front, so the legacy fixed 64k-token pool
    (4096 blocks) that is harmless on a 128 GB box is several GB of dead-weight KV on a
    low-RAM box — enough to tip generation into swap on top of the model weights and the
    resident embeddings model. Scale it to unified RAM. A pool smaller than the prompt only
    caches less of the prefix (a perf trade), it never breaks correctness. Override with
    MUGI_MLX_APC_NUM_BLOCKS or APC_NUM_BLOCKS. Each block = 16 tokens (4096 = 64k tokens).
    """
    ram_gb = _get_system_ram_gb()
    if not isinstance(ram_gb, int):
        return 2048  # ~32k tokens: safe middle ground when RAM can't be determined
    if ram_gb <= 16:
        return 1024  # ~16k tokens
    if ram_gb <= 24:
        return 1536  # ~24k tokens
    if ram_gb <= 32:
        return 2048  # ~32k tokens
    if ram_gb <= 48:
        return 3072  # ~48k tokens
    return 4096      # ~64k tokens (legacy default; only for >48 GB machines)


# ---------------------------------------------------------------------------
# Model download
# ---------------------------------------------------------------------------

_download_jobs: dict[str, dict] = {}
_download_lock = Lock()

# Hugging Face model repo id (namespace/name), no path traversal.
_HF_REPO_ID_RE = re.compile(r"^[\w.\-]+/[\w.\-]+$")


def _repo_dirname(hf_repo: str) -> str:
    parts = [p for p in (hf_repo or "").strip().split("/") if p]
    if len(parts) >= 2:
        return f"{parts[0]}__{parts[1]}"
    return re.sub(r"[^\w.\-]+", "_", (hf_repo or "model").strip())


def _hf_cache_snapshots_dir(hf_repo: str) -> str:
    parts = [p for p in (hf_repo or "").strip().split("/") if p]
    if len(parts) < 2:
        return ""
    return os.path.expanduser(os.path.join("~/.cache/huggingface/hub", f"models--{parts[0]}--{parts[1]}", "snapshots"))


def _validate_model_dir(path: str) -> tuple[bool, str]:
    if not os.path.isdir(path):
        return False, "Model path is not a directory."
    if not os.path.isfile(os.path.join(path, "config.json")):
        return False, "config.json not found."
    if not _is_mlx_model_dir(path):
        return False, "Model directory is not MLX format (expected safetensors and no GGUF)."
    return True, ""


def _resolve_snapshot_path(hf_repo: str, revision: str | None = None) -> str:
    snapshots = _hf_cache_snapshots_dir(hf_repo)
    if not snapshots or not os.path.isdir(snapshots):
        return ""
    wanted = (revision or "").strip()
    if wanted:
        candidate = os.path.join(snapshots, wanted)
        if os.path.isdir(candidate):
            return candidate
    newest = ""
    newest_mtime = -1.0
    try:
        for entry in os.scandir(snapshots):
            if not entry.is_dir():
                continue
            try:
                mtime = entry.stat().st_mtime
            except OSError:
                mtime = 0.0
            if mtime >= newest_mtime:
                newest_mtime = mtime
                newest = entry.path
    except OSError:
        return ""
    return newest


def _finalize_downloaded_model(hf_repo: str, revision: str | None = None) -> tuple[bool, str, str]:
    """Validate and canonicalize a downloaded model path under ~/.mugi/mlx/models."""
    snapshot = _resolve_snapshot_path(hf_repo, revision)
    if not snapshot:
        return False, "Downloaded snapshot path could not be resolved.", ""
    ok, msg = _validate_model_dir(snapshot)
    if not ok:
        return False, msg, ""
    os.makedirs(MLX_CANONICAL_MODELS_DIR, exist_ok=True)
    canonical = os.path.join(MLX_CANONICAL_MODELS_DIR, _repo_dirname(hf_repo))
    tmp_link = canonical + ".tmp"
    try:
        if os.path.lexists(tmp_link):
            if os.path.isdir(tmp_link) and not os.path.islink(tmp_link):
                shutil.rmtree(tmp_link)
            else:
                os.remove(tmp_link)
        os.symlink(snapshot, tmp_link)
        if os.path.lexists(canonical):
            if os.path.isdir(canonical) and not os.path.islink(canonical):
                shutil.rmtree(canonical)
            else:
                os.remove(canonical)
        os.replace(tmp_link, canonical)
        meta = {
            "hf_repo": hf_repo,
            "revision": (revision or ""),
            "snapshot_path": snapshot,
            "finalized_at": int(time.time()),
        }
        with open(os.path.join(canonical, ".mugi_model_meta.json"), "w", encoding="utf-8") as f:
            json.dump(meta, f, indent=2)
    except Exception as exc:
        try:
            if os.path.lexists(tmp_link):
                os.remove(tmp_link)
        except Exception:
            pass
        return False, f"Model finalization failed: {exc}", ""
    return True, "Model finalized.", canonical

# Child process (MLX venv Python): snapshot_download with tqdm hooks → JSON lines on stdout.
# Note: HF_HUB_ENABLE_HF_TRANSFER forces a sequential path without thread_map/tqdm; we clear it here
# so tqdm_class runs (slightly slower than hf_transfer, but progress is visible).
_MLX_HF_DOWNLOAD_WORKER = r'''import json
import os
import sys
import threading
import time

REPO_ID = sys.argv[1]
REVISION = sys.argv[2] if len(sys.argv) > 2 and (sys.argv[2] or "").strip() else None


def _emit(progress, downloaded, total, message=""):
    print(
        json.dumps(
            {
                "progress": round(progress, 6),
                "downloaded_bytes": int(downloaded),
                "total_bytes": int(total),
                "message": message,
            }
        ),
        flush=True,
    )


_last_emit = [0.0, -1.0]


def _emit_throttled(progress, downloaded, total, message=""):
    now = time.monotonic()
    if now - _last_emit[0] < 0.2 and abs(progress - _last_emit[1]) < 0.002:
        return
    _last_emit[0] = now
    _last_emit[1] = progress
    _emit(progress, downloaded, total, message)


def _estimate_repo_bytes():
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        total = 0
        if not hasattr(api, "list_repo_tree"):
            return 0
        for entry in api.list_repo_tree(repo_id=REPO_ID, repo_type="model", recursive=True):
            sz = getattr(entry, "size", None)
            if sz is not None:
                try:
                    total += int(sz)
                except (TypeError, ValueError):
                    pass
        return total
    except Exception:
        return 0


REPO_TOTAL = _estimate_repo_bytes()


class _Coord:
    def __init__(self):
        self.lock = threading.Lock()
        self.active = {}
        self.completed_before = 0
        self.repo_total = REPO_TOTAL

    def _progress_unlocked(self):
        active_sum = sum(b.n for b in self.active.values())
        done_files = self.completed_before + active_sum
        file_total = 0
        for b in self.active.values():
            file_total = max(file_total, b.total)
        if file_total <= 0:
            prog = 0.0
        else:
            prog = min(1.0, done_files / file_total)
        rt = self.repo_total
        est_bytes = int(rt * prog) if rt > 0 else 0
        return prog, est_bytes, rt, done_files, file_total

    def emit(self):
        with self.lock:
            prog, est_bytes, rt, done_files, ft = self._progress_unlocked()
            if rt > 0:
                msg = f"Downloading {REPO_ID}… ({est_bytes / (1024 ** 3):.2f} / {rt / (1024 ** 3):.2f} GB)"
            elif ft > 0:
                msg = f"Downloading {REPO_ID}… (file {done_files} of {ft})"
            else:
                msg = f"Downloading {REPO_ID}…"
        _emit_throttled(prog, est_bytes, rt, msg)


COORD = _Coord()


class _Tqdm:
    """Enough of tqdm's API for huggingface_hub.snapshot_download(thread_map(..., tqdm_class=...)).

    tqdm.contrib.concurrent.thread_map uses ensure_lock(), which does ``del tqdm_class._lock`` on exit
    if no lock existed initially. A tqdm-compatible class must use ``_lock`` (not a private alias), or
    cleanup raises AttributeError.
    """

    _lock = None

    @classmethod
    def get_lock(cls):
        if not hasattr(cls, "_lock") or cls._lock is None:
            cls._lock = threading.RLock()
        return cls._lock

    @classmethod
    def set_lock(cls, lock=None):
        cls._lock = lock

    def __init__(self, *args, **kwargs):
        self._wrapped = args[0] if args else None
        self.total = int(kwargs.get("total") or 0)
        self.n = int(kwargs.get("initial") or 0)
        self._it = None
        self._id = id(self)
        self._closed = False
        with COORD.lock:
            COORD.active[self._id] = self
        COORD.emit()

    def __iter__(self):
        if self._wrapped is None:
            return iter(())
        self._it = iter(self._wrapped)
        return self

    def __next__(self):
        if self._it is None:
            self._it = iter(self._wrapped) if self._wrapped is not None else iter(())
        try:
            item = next(self._it)
        except StopIteration:
            self.close()
            raise
        self.update(1)
        return item

    def update(self, n=1):
        self.n += n
        COORD.emit()

    def close(self):
        if self._closed:
            return
        self._closed = True
        with COORD.lock:
            COORD.active.pop(self._id, None)
            COORD.completed_before += self.n
        COORD.emit()

    def __enter__(self):
        return self

    def __exit__(self, *a):
        self.close()
        return False

    def refresh(self):
        pass

    def set_description(self, *args, **kwargs):
        pass

    def set_postfix(self, *args, **kwargs):
        pass

    def set_postfix_str(self, *args, **kwargs):
        pass

    def clear(self, *args, **kwargs):
        pass

    def unpause(self, *args, **kwargs):
        pass


def main():
    os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None)
    try:
        from huggingface_hub import snapshot_download

        kwargs = {"repo_type": "model"}
        if REVISION:
            kwargs["revision"] = REVISION
        try:
            snapshot_download(REPO_ID, tqdm_class=_Tqdm, **kwargs)
        except TypeError:
            snapshot_download(REPO_ID, **kwargs)
    except Exception as e:
        print(json.dumps({"event": "error", "message": str(e)[:800]}), flush=True)
        sys.exit(1)

    with COORD.lock:
        for b in list(COORD.active.values()):
            COORD.completed_before += b.n
        COORD.active.clear()
        done_files = COORD.completed_before
        rt = COORD.repo_total
    est_bytes = rt if rt > 0 else 0
    _emit(1.0, est_bytes if rt > 0 else done_files, rt if rt > 0 else 0, f"{REPO_ID} downloaded.")
    print(json.dumps({"event": "done"}), flush=True)


if __name__ == "__main__":
    main()
'''


def _mlx_download_merge_worker_json(status: dict, obj: dict) -> str | None:
    """Apply one decoded JSON object from the MLX HF download worker stdout.

    Updates ``status`` in place. Returns ``"done"``, ``"error"``, or None.
    """
    ev = obj.get("event")
    if ev == "error":
        msg = obj.get("message")
        status["status"] = "failed"
        status["message"] = (str(msg) if msg else "Download failed.")[:500]
        status["progress"] = status.get("progress", 0.0)
        return "error"
    if ev == "done":
        status["status"] = "completed"
        status["progress"] = 1.0
        rid = str(status.get("hf_repo") or "")
        status["message"] = f"{rid} downloaded." if rid else "Download complete."
        tb = int(status.get("total_bytes") or 0)
        db = int(status.get("downloaded_bytes") or 0)
        status["downloaded_bytes"] = tb if tb > 0 else db
        return "done"
    pr = obj.get("progress")
    if isinstance(pr, (int, float)):
        status["progress"] = float(min(1.0, max(0.0, pr)))
    db = obj.get("downloaded_bytes")
    if isinstance(db, int):
        status["downloaded_bytes"] = db
    elif isinstance(db, float):
        status["downloaded_bytes"] = int(db)
    tb = obj.get("total_bytes")
    if isinstance(tb, int):
        status["total_bytes"] = tb
    elif isinstance(tb, float):
        status["total_bytes"] = int(tb)
    msg = obj.get("message")
    if isinstance(msg, str) and msg.strip():
        status["message"] = msg.strip()[:500]
    return None


def mlx_download_model(hf_repo: str, revision: str | None = None) -> dict:
    """Start downloading an MLX model from HuggingFace.
    Returns {"ok": bool, "message": str, "job_id": str?}.
    """
    if not mlx_venv_exists():
        return {"ok": False, "message": "MLX venv not installed."}

    hf_repo = (hf_repo or "").strip()
    if not _HF_REPO_ID_RE.match(hf_repo):
        return {"ok": False, "message": "Invalid Hugging Face repo id."}

    import uuid
    from threading import Thread

    job_id = str(uuid.uuid4())
    status: dict = {
        "status": "downloading",
        "progress": 0.0,
        "message": f"Downloading {hf_repo}...",
        "downloaded_bytes": 0,
        "total_bytes": 0,
        "hf_repo": hf_repo,
        "revision": (revision or "").strip(),
    }
    with _download_lock:
        _download_jobs[job_id] = status

    py = _mlx_python()
    if not os.path.isfile(py):
        with _download_lock:
            _download_jobs.pop(job_id, None)
        return {"ok": False, "message": "MLX venv Python not found."}

    def _do_download():
        script_path = None
        proc: subprocess.Popen | None = None
        try:
            fd, script_path = tempfile.mkstemp(prefix="mugi_mlx_dl_", suffix=".py")
            with os.fdopen(fd, "w", encoding="utf-8") as wf:
                wf.write(_MLX_HF_DOWNLOAD_WORKER)
            cmd = [py, script_path, hf_repo]
            rev = (revision or "").strip()
            if rev:
                cmd.append(rev)
            env = {**os.environ, "HF_HUB_ENABLE_HF_TRANSFER": "1"}
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                env=env,
            )
            err_chunks: list[str] = []

            def _pump_stderr() -> None:
                try:
                    if proc is not None and proc.stderr:
                        while True:
                            chunk = proc.stderr.read(4096)
                            if not chunk:
                                break
                            err_chunks.append(chunk)
                except Exception:
                    pass

            from threading import Thread as _Thr

            _Thr(target=_pump_stderr, daemon=True).start()

            if proc.stdout is None:
                raise RuntimeError("stdout not piped")
            assert proc is not None
            for line in iter(proc.stdout.readline, ""):
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                with _download_lock:
                    st = _download_jobs.get(job_id)
                    if not st:
                        break
                    ev = _mlx_download_merge_worker_json(st, obj)
                if ev == "error":
                    proc.wait(timeout=5)
                    return
                if ev == "done":
                    proc.wait(timeout=5)
                    with _download_lock:
                        st = _download_jobs.get(job_id)
                        if st:
                            ok_final, final_msg, canonical_path = _finalize_downloaded_model(hf_repo, revision=revision)
                            if ok_final:
                                st["status"] = "completed"
                                st["progress"] = 1.0
                                st["message"] = f"{hf_repo} downloaded and finalized."
                                st["final_model_path"] = canonical_path
                            else:
                                st["status"] = "failed"
                                st["message"] = final_msg[:500]
                    return

            rc = proc.wait(timeout=60)
            with _download_lock:
                st = _download_jobs.get(job_id)
                if not st:
                    return
                if st.get("status") == "completed":
                    ok_final, final_msg, canonical_path = _finalize_downloaded_model(hf_repo, revision=revision)
                    if ok_final:
                        st["message"] = f"{hf_repo} downloaded and finalized."
                        st["final_model_path"] = canonical_path
                    else:
                        st["status"] = "failed"
                        st["message"] = final_msg[:500]
                    return
                if rc == 0:
                    ok_final, final_msg, canonical_path = _finalize_downloaded_model(hf_repo, revision=revision)
                    if ok_final:
                        st["status"] = "completed"
                        st["progress"] = 1.0
                        st["message"] = f"{hf_repo} downloaded and finalized."
                        st["final_model_path"] = canonical_path
                    else:
                        st["status"] = "failed"
                        st["message"] = final_msg[:500]
                else:
                    st["status"] = "failed"
                    tail = "".join(err_chunks).strip()
                    st["message"] = (tail or f"Download failed (exit {rc}).")[:500]
        except subprocess.TimeoutExpired:
            with _download_lock:
                st = _download_jobs.get(job_id)
                if st and st.get("status") == "downloading":
                    st["status"] = "failed"
                    st["message"] = "Download timed out."
            try:
                if proc is not None:
                    proc.kill()
            except Exception:
                pass
        except Exception as e:
            with _download_lock:
                st = _download_jobs.get(job_id)
                if st and st.get("status") == "downloading":
                    st["status"] = "failed"
                    st["message"] = str(e)[:500]
        finally:
            if script_path:
                try:
                    os.remove(script_path)
                except OSError:
                    pass

    Thread(target=_do_download, daemon=True).start()
    return {"ok": True, "message": "Download started.", "job_id": job_id}


def mlx_download_status(job_id: str) -> dict:
    """Return download progress for a job."""
    with _download_lock:
        job = _download_jobs.get(job_id)
        if not job:
            return {"status": "unknown", "message": "Job not found."}
        return dict(job)


# ---------------------------------------------------------------------------
# Adapter management
# ---------------------------------------------------------------------------


def mlx_list_adapters() -> list[dict]:
    """Return list of saved adapters.
    Each: {"name": str, "path": str, "created": str}.
    """
    if not os.path.isdir(MLX_ADAPTERS_DIR):
        return []
    adapters = []
    try:
        for entry in os.scandir(MLX_ADAPTERS_DIR):
            if not entry.is_dir():
                continue
            config_path = os.path.join(entry.path, "adapter_config.json")
            weights_path = os.path.join(entry.path, "adapters.safetensors")
            if not (os.path.isfile(config_path) or os.path.isfile(weights_path)):
                continue
            created = ""
            try:
                st = os.stat(entry.path)
                created = time.strftime("%Y-%m-%d %H:%M", time.localtime(st.st_ctime))
            except OSError:
                pass
            adapters.append({
                "name": entry.name,
                "path": entry.path,
                "created": created,
            })
    except OSError:
        pass
    adapters.sort(key=lambda a: a["name"])
    return adapters


def mlx_delete_adapter(name: str) -> tuple[bool, str]:
    """Delete a named adapter directory. Returns (success, message)."""
    if not name or ".." in name or "/" in name:
        return False, "Invalid adapter name."
    path = os.path.join(MLX_ADAPTERS_DIR, name)
    if not os.path.isdir(path):
        return False, "Adapter not found."
    try:
        shutil.rmtree(path)
        return True, f"Adapter '{name}' deleted."
    except Exception as e:
        return False, str(e)


def mlx_get_active_adapter() -> str:
    """Read mlx_active_adapter from app_prefs.json. Returns adapter name or empty string."""
    try:
        from utils.app_prefs_store import load_app_prefs_cached
        prefs = load_app_prefs_cached()
        return (prefs.get("mlx_active_adapter") or "").strip()
    except Exception:
        pass
    return ""
