"""Embedded utility model inference via mlx-vlm.

This module provides a lightweight, in-process completion API for repetitive
background tasks (summarization/classification/rewrite). It is intentionally
best-effort: callers should always keep a fallback path to the main model.

Default weights are mlx-community/gemma-4-E2B-it-qat-8bit (Gemma 4 E2B QAT, ~2B-effective;
see persistence.DEFAULT_UTILITY_MODEL_REPO) — chosen because it installs cleanly via
mlx-vlm and does not require the full Xcode + Metal Toolchain + CMake. Users who opt
into prism-ml/Bonsai-8B-mlx-1bit need the Prism MLX fork, which has no prebuilt wheels
and compiles from source. Uses ~/.mugi/mlx-utility-venv (mlx-vlm, with optional Prism
MLX for Bonsai), separate from the main ~/.mugi/mlx-venv used for the mlx-vlm server.
"""
from __future__ import annotations

import atexit
import gc
import hashlib
import json
import math
import os
import re
import select
import site
import subprocess
import sys
import textwrap
import threading
import time
from typing import Callable

from app_mlx import (
    MLX_UTILITY_VENV_DIR,
    ensure_utility_mlx_venv,
    utility_mlx_prism_ready,
    utility_mlx_venv_exists,
)
from core.config import get_utility_config
from persistence import DEFAULT_UTILITY_MODEL_REPO

_LOAD_LOCK = threading.Lock()
_MODEL = None
_PROCESSOR = None
_CONFIG = None
_UTILITY_PKGS_READY = False
_LOAD_BACKEND = ""
_LAST_LOAD_ERROR = ""

_ACTIVITY_SOURCE = "utility_model"
_ACTIVITY_OP_DEFS: dict[str, dict[str, str]] = {
    "context_summarization": {
        "started": "Utility model: summarizing conversation context...",
        "completed": "Utility model: context summary complete",
        "error": "Utility model: context summarization failed",
    },
    "knowledge_digest": {
        "started": "Utility model: digesting knowledge source...",
        "completed": "Utility model: knowledge digest complete",
        "error": "Utility model: knowledge digest failed",
    },
    "memory_refresh": {
        "started": "Utility model: refreshing stale memory entries...",
        "completed": "Utility model: memory refresh analysis complete",
        "error": "Utility model: memory refresh analysis failed",
    },
    "session_search_summary": {
        "started": "Utility model: summarizing past session match...",
        "completed": "Utility model: session summary complete",
        "error": "Utility model: session summary failed",
    },
    "session_search_embedding": {
        "started": "Utility model: embedding session search text...",
        "completed": "Utility model: embedding complete",
        "error": "Utility model: embedding failed",
    },
    "ocr_text_correction": {
        "started": "Utility model: correcting OCR text...",
        "completed": "Utility model: OCR correction complete",
        "error": "Utility model: OCR correction failed",
    },
    "handoff_critic": {
        "started": "Utility model: reviewing handoff consistency...",
        "completed": "Utility model: handoff review complete",
        "error": "Utility model: handoff review failed",
    },
    "completion_critic": {
        "started": "Utility model: reviewing final reply vs tool evidence...",
        "completed": "Utility model: completion review done",
        "error": "Utility model: completion review failed",
    },
    "conversation_title": {
        "started": "Utility model: generating conversation title...",
        "completed": "Utility model: conversation title ready",
        "error": "Utility model: conversation title failed",
    },
    "ledger_reflection": {
        "started": "Utility model: drafting weekly self-review...",
        "completed": "Utility model: weekly self-review draft ready",
        "error": "Utility model: weekly self-review failed",
    },
    "thread_distillation": {
        "started": "Utility model: distilling thread facts...",
        "completed": "Utility model: thread distillation complete",
        "error": "Utility model: thread distillation failed",
    },
    "web_act_reflection": {
        "started": "Utility model: summarizing web automation run...",
        "completed": "Utility model: web automation summary ready",
        "error": "Utility model: web automation summary failed",
    },
    "generic": {
        "started": "Utility model: running background task...",
        "completed": "Utility model: background task complete",
        "error": "Utility model: background task failed",
    },
}


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


def _venv_site_packages() -> list[str]:
    """Return plausible site-packages paths for the utility MLX venv."""
    paths: list[str] = []
    # Prefer discovering actual interpreter dirs inside the utility venv, since
    # this process Python version can differ from the utility venv version.
    lib_root = os.path.join(MLX_UTILITY_VENV_DIR, "lib")
    try:
        for entry in sorted(os.listdir(lib_root)):
            if not entry.startswith("python"):
                continue
            paths.append(os.path.join(lib_root, entry, "site-packages"))
    except Exception:
        pass
    # Back-compat guess using current runtime version.
    ver = f"python{sys.version_info.major}.{sys.version_info.minor}"
    paths.append(os.path.join(MLX_UTILITY_VENV_DIR, "lib", ver, "site-packages"))
    paths.extend(site.getsitepackages([MLX_UTILITY_VENV_DIR]))
    # De-duplicate while preserving order.
    seen: set[str] = set()
    out: list[str] = []
    for p in paths:
        if p and p not in seen and os.path.isdir(p):
            out.append(p)
            seen.add(p)
    return out


def _ensure_utility_packages_importable() -> bool:
    """Add utility MLX venv site-packages to sys.path."""
    global _UTILITY_PKGS_READY
    if _UTILITY_PKGS_READY:
        return True

    for sp in _venv_site_packages():
        if sp not in sys.path:
            sys.path.insert(0, sp)

    _UTILITY_PKGS_READY = True
    return True


def _get_model_repo() -> str:
    cfg = get_utility_config()
    repo = str(cfg.get("model_repo") or "").strip()
    return repo or DEFAULT_UTILITY_MODEL_REPO


def _local_model_path_if_any(model_ref: str) -> str | None:
    """Return expanded local path when model_ref points to local filesystem."""
    ref = str(model_ref or "").strip()
    if not ref:
        return None
    expanded = os.path.abspath(os.path.expanduser(ref))
    if os.path.exists(expanded):
        return expanded
    if ref.startswith(("/", "~", "./", "../")):
        return expanded
    return None


def _enabled() -> bool:
    cfg = get_utility_config()
    return bool(cfg.get("enabled", True))


def _activity_definition(operation: str | None) -> tuple[str, dict[str, str]]:
    op = str(operation or "").strip().lower() or "generic"
    return op, _ACTIVITY_OP_DEFS.get(op, _ACTIVITY_OP_DEFS["generic"])


def _emit_activity(
    state: str,
    summary: str,
    *,
    operation: str,
    job_id: str,
    thread_id: str | None = None,
) -> None:
    try:
        from activity_bus import emit_activity_enriched
        event = {
            "source": _ACTIVITY_SOURCE,
            "state": state,
            "summary": summary,
            "job_id": job_id,
            "operation": operation,
        }
        if thread_id:
            event["thread_id"] = thread_id
        emit_activity_enriched(event, tags=["background"])
    except Exception:
        pass


def _flatten_content(content) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for p in content:
            if not isinstance(p, dict):
                continue
            if p.get("type") == "text":
                txt = p.get("text")
                if isinstance(txt, str) and txt.strip():
                    parts.append(txt.strip())
        return "\n".join(parts)
    return ""


def _messages_to_prompt(messages: list[dict]) -> str:
    parts: list[str] = []
    for m in messages:
        if not isinstance(m, dict):
            continue
        role = str(m.get("role") or "user").upper()
        content = _flatten_content(m.get("content"))
        if content.strip():
            parts.append(f"{role}: {content}")
    return "\n\n".join(parts).strip()


def _apply_utility_chat_template(processor, messages: list[dict]) -> str | None:
    """Format messages for instruction-tuned utility models (Gemma 4, Qwen, etc.)."""
    if not hasattr(processor, "apply_chat_template"):
        return None
    variants = (
        {"tokenize": False, "add_generation_prompt": True, "enable_thinking": False},
        {
            "tokenize": False,
            "add_generation_prompt": True,
            "chat_template_kwargs": {"enable_thinking": False},
        },
        {"tokenize": False, "add_generation_prompt": True},
    )
    for kwargs in variants:
        try:
            return processor.apply_chat_template(messages, **kwargs)
        except TypeError:
            continue
        except Exception:
            break
    return None


def _thinking_inner_fallback(text: str) -> str:
    """When Gemma emits only a thinking block, salvage its inner text."""
    raw = str(text or "")
    if not raw.strip():
        return ""
    for pattern in (
        r"<\|channel>thought([\s\S]*?)<channel\|>",
        r"<think>([\s\S]*?)</think>",
    ):
        match = re.search(pattern, raw, flags=re.IGNORECASE)
        if match:
            inner = match.group(1).strip()
            if inner:
                return inner
    for open_tag in ("<|channel>thought", "<think>"):
        idx = raw.lower().find(open_tag.lower())
        if idx != -1:
            tail = raw[idx + len(open_tag):].strip()
            if tail:
                return tail
    return ""


def _normalize_utility_output(raw: str) -> str:
    """Strip private reasoning markup; fall back to thinking inner text when needed."""
    from utils.llm_output_sanitize import strip_think_tags

    cleaned = strip_think_tags(str(raw or "")).strip()
    if cleaned:
        return cleaned
    return _thinking_inner_fallback(str(raw or "")).strip()


def _strip_thinking_blocks(text: str) -> str:
    """Backward-compatible alias for callers/tests."""
    return _normalize_utility_output(text)


def _decode_generate_result(result) -> str:
    if result is None:
        return ""
    # mlx-vlm 0.6.x returns a GenerationResult dataclass whose ``str()`` is the full repr
    # (e.g. "GenerationResult(text='pong', token=106, logprobs=...)"); the actual completion
    # lives on ``.text``. mlx-lm still returns a plain string. Prefer ``.text`` so the vlm
    # backend doesn't hand callers the repr instead of the answer.
    text_attr = getattr(result, "text", None)
    if isinstance(text_attr, str):
        return text_attr
    if isinstance(result, tuple):
        return str(result[0] or "")
    return str(result or "")


def _utility_retry_messages(messages: list[dict]) -> list[dict]:
    return [
        {
            "role": "system",
            "content": (
                "Respond with only the final answer. No thinking tags, channel markup, "
                "or reasoning preamble."
            ),
        },
        *messages,
    ]


# ---------------------------------------------------------------------------
# Persistent subprocess worker — runs in the utility venv's own Python so
# compiled mlx extensions always match the interpreter ABI.
# ---------------------------------------------------------------------------

# Shared monkeypatch used by BOTH the in-process loader (``_ensure_loaded``) and the
# subprocess worker script below. mlx-vlm 0.6.2 skips its weight sanitize for mlx-format
# checkpoints (``utils.load_model``: ``if not is_mlx_format``), but the mlx-format Gemma-4
# E2B conversions ship redundant per-layer k_proj/v_proj/k_norm for their KV-shared layers
# (config ``num_kv_shared_layers``). gemma4's model class builds those layers as
# ``kv_shared_only`` (no own projections), so the strict ``load_weights`` then raises
# "Received N parameters not in model" and the utility model is reported unavailable. We
# strip exactly the weights gemma4's own ``LanguageModel.sanitize`` would have dropped,
# restoring a clean load. Idempotent and fully defensive (never raises). The single source
# of truth lives here as text so the separate-interpreter worker can ``-c`` it verbatim and
# the in-process path can ``exec`` the identical logic (no divergence).
_GEMMA4_KV_SHARE_PATCH = '''\
def _apply_gemma4_kv_share_patch():
    try:
        from mlx_vlm.models.gemma4 import gemma4 as _g4
    except Exception:
        return
    _orig = _g4.Model.load_weights
    if getattr(_orig, "_mugi_kv_share_patched", False):
        return

    def _patched(self, weights, strict=True):
        lm = getattr(self, "language_model", None)
        pred = getattr(lm, "_is_unused_shared_kv_weight", None)
        if callable(pred):
            try:
                if isinstance(weights, dict):
                    weights = {k: v for k, v in weights.items() if not pred(k)}
                elif isinstance(weights, (list, tuple)):
                    weights = [(k, v) for (k, v) in weights if not pred(k)]
            except Exception:
                pass
        return _orig(self, weights, strict=strict)

    _patched._mugi_kv_share_patched = True
    _g4.Model.load_weights = _patched


_apply_gemma4_kv_share_patch()
'''


_WORKER_SCRIPT = _GEMMA4_KV_SHARE_PATCH + textwrap.dedent("""\
    import sys, json, os

    repo = sys.argv[1]

    # Try mlx_lm first (non-VLM models like Bonsai), then mlx_vlm.
    backend = None
    model = tokenizer = processor = config = None
    try:
        from mlx_lm import load as load_lm, generate as generate_lm
        model, tokenizer = load_lm(repo)
        backend = "lm"
    except Exception as e1:
        try:
            from mlx_vlm import load as load_vlm, generate as generate_vlm
            from mlx_vlm.utils import load_config
            model, processor = load_vlm(repo)
            config = load_config(repo)
            backend = "vlm"
        except Exception as e2:
            print(json.dumps({"ready": False, "error": f"lm: {e1} | vlm: {e2}"}),
                  flush=True)
            sys.exit(1)

    print(json.dumps({"ready": True, "backend": backend}), flush=True)

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            messages = req.get("messages")
            prompt_text = req.get("prompt", "")
            max_tokens = req.get("max_tokens", 512)
            temperature = req.get("temperature", 0.3)

            if backend == "lm":
                # Apply chat template when we have structured messages and the
                # tokenizer supports it — disable thinking for Gemma 4 utility work.
                if messages and hasattr(tokenizer, "apply_chat_template"):
                    formatted = None
                    for kwargs in (
                        {"tokenize": False, "add_generation_prompt": True, "enable_thinking": False},
                        {
                            "tokenize": False,
                            "add_generation_prompt": True,
                            "chat_template_kwargs": {"enable_thinking": False},
                        },
                        {"tokenize": False, "add_generation_prompt": True},
                    ):
                        try:
                            formatted = tokenizer.apply_chat_template(messages, **kwargs)
                            break
                        except TypeError:
                            continue
                        except Exception:
                            break
                    if formatted:
                        prompt_text = formatted
                try:
                    result = generate_lm(
                        model, tokenizer, prompt=prompt_text,
                        max_tokens=max_tokens, temperature=temperature, verbose=False,
                    )
                except TypeError:
                    result = generate_lm(
                        model, tokenizer, prompt=prompt_text,
                        max_tokens=max_tokens, verbose=False,
                    )
            else:
                from mlx_vlm.prompt_utils import apply_chat_template
                try:
                    fmt = apply_chat_template(processor, config, prompt_text,
                                              num_images=0, enable_thinking=False)
                except TypeError:
                    fmt = apply_chat_template(processor, config, prompt_text,
                                              num_images=0)
                result = generate_vlm(model, processor, fmt,
                                      max_tokens=max_tokens,
                                      temperature=temperature, verbose=False)

            # mlx-vlm 0.6.x returns a GenerationResult (.text holds the completion;
            # str() would yield the dataclass repr); mlx-lm returns a plain string.
            _t = getattr(result, "text", None)
            if isinstance(_t, str):
                text = _t
            elif isinstance(result, tuple):
                text = str(result[0] or "")
            else:
                text = str(result or "")
            print(json.dumps({"ok": True, "text": text}), flush=True)
        except Exception as e:
            print(json.dumps({"ok": False, "error": str(e)[:300]}), flush=True)
""")


class _SubprocessWorker:
    """Persistent utility-venv subprocess that loads the model once."""

    def __init__(self) -> None:
        self._proc: subprocess.Popen | None = None
        self._ready = False
        self._backend = ""

    def start(self, repo: str, *, timeout: float = 120) -> bool:
        py = _mlx_python()
        if not os.path.isfile(py):
            return False
        try:
            self._proc = subprocess.Popen(
                [py, "-c", _WORKER_SCRIPT, repo],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                env={**os.environ, "HF_HUB_ENABLE_HF_TRANSFER": "1"},
            )
            ready, _, _ = select.select([self._proc.stdout], [], [], timeout)
            if not ready:
                self._proc.kill()
                return False
            line = self._proc.stdout.readline()  # type: ignore[union-attr]
            if not line:
                self._proc.kill()
                return False
            msg = json.loads(line.strip())
            if msg.get("ready"):
                self._ready = True
                self._backend = msg.get("backend", "")
                return True
            self._proc.kill()
            return False
        except Exception:
            if self._proc:
                try:
                    self._proc.kill()
                except Exception:
                    pass
            self._proc = None
            return False

    @property
    def alive(self) -> bool:
        return (self._proc is not None
                and self._proc.poll() is None
                and self._ready)

    def complete(self, prompt: str, max_tokens: int, temperature: float,
                 *, messages: list[dict] | None = None,
                 timeout: float = 180) -> str | None:
        if not self.alive:
            return None
        try:
            req = json.dumps({
                "prompt": prompt,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
            })
            self._proc.stdin.write(req + "\n")  # type: ignore[union-attr]
            self._proc.stdin.flush()  # type: ignore[union-attr]
            ready, _, _ = select.select([self._proc.stdout], [], [], timeout)
            if not ready:
                return None
            line = self._proc.stdout.readline()  # type: ignore[union-attr]
            if not line:
                return None
            resp = json.loads(line.strip())
            if resp.get("ok"):
                return str(resp.get("text") or "").strip() or None
            return None
        except Exception:
            return None

    def stop(self) -> None:
        proc = self._proc
        self._proc = None
        self._ready = False
        if proc and proc.poll() is None:
            try:
                proc.stdin.close()  # type: ignore[union-attr]
                proc.wait(timeout=5)
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass


_WORKER: _SubprocessWorker | None = None
_WORKER_LOCK = threading.Lock()


def _stop_subprocess_worker() -> None:
    global _WORKER
    with _WORKER_LOCK:
        if _WORKER is not None:
            _WORKER.stop()
            _WORKER = None


atexit.register(_stop_subprocess_worker)


def _subprocess_utility_complete(
    messages: list[dict],
    max_tokens: int,
    temperature: float,
) -> str | None:
    """Complete via persistent subprocess in the utility venv (ABI-safe)."""
    global _WORKER
    prompt = _messages_to_prompt(messages)
    if not prompt:
        return None
    with _WORKER_LOCK:
        if _WORKER is None or not _WORKER.alive:
            repo = _get_model_repo()
            worker = _SubprocessWorker()
            if not worker.start(repo):
                return None
            _WORKER = worker
        text = _WORKER.complete(prompt, max_tokens, temperature, messages=messages)
        normalized = _normalize_utility_output(text) if text else ""
        if normalized:
            return normalized
    retry_messages = _utility_retry_messages(messages)
    retry_prompt = _messages_to_prompt(retry_messages)
    if retry_prompt and retry_prompt != prompt:
        with _WORKER_LOCK:
            if _WORKER is not None and _WORKER.alive:
                text = _WORKER.complete(
                    retry_prompt,
                    max_tokens,
                    temperature,
                    messages=retry_messages,
                )
                if text:
                    normalized = _normalize_utility_output(text)
                    if normalized:
                        return normalized
    return None


def _ensure_loaded() -> bool:
    global _MODEL, _PROCESSOR, _CONFIG, _LOAD_BACKEND, _LAST_LOAD_ERROR
    if _MODEL is not None and _PROCESSOR is not None and (_CONFIG is not None or _LOAD_BACKEND == "lm"):
        return True
    ok, msg = ensure_utility_mlx_venv()
    if not ok:
        _LAST_LOAD_ERROR = str(msg or "utility mlx venv setup failed")
        return False
    if not _ensure_utility_packages_importable():
        _LAST_LOAD_ERROR = "utility site-packages unavailable"
        return False

    # mlx-format Gemma-4 (E2B) checkpoints carry redundant KV-shared weights that
    # mlx-vlm's strict loader rejects; apply the shared-KV drop before loading. Same
    # trusted internal source as the subprocess worker (see _GEMMA4_KV_SHARE_PATCH).
    try:
        exec(_GEMMA4_KV_SHARE_PATCH, {})  # noqa: S102 - trusted module-level constant
    except Exception:
        pass

    repo = _get_model_repo()
    lower_repo = repo.lower()
    prefer_vlm = any(tag in lower_repo for tag in ("-vl", "vl-", "vision", "4v", "smolvlm"))

    def _try_load_lm() -> tuple[bool, object]:
        try:
            from mlx_lm import load as load_lm
            model, processor = load_lm(repo)
            return True, (model, processor)
        except Exception as exc:
            return False, str(exc).strip()[:220]

    def _try_load_vlm() -> tuple[bool, object]:
        try:
            from mlx_vlm import load as load_vlm
            from mlx_vlm.utils import load_config
            model, processor = load_vlm(repo)
            cfg = load_config(repo)
            return True, (model, processor, cfg)
        except Exception as exc:
            return False, str(exc).strip()[:220]

    lm_err = ""
    vlm_err = ""
    if prefer_vlm:
        ok_vlm, vlm_payload = _try_load_vlm()
        if ok_vlm:
            _MODEL, _PROCESSOR, _CONFIG = vlm_payload  # type: ignore[misc]
            _LOAD_BACKEND = "vlm"
            _LAST_LOAD_ERROR = ""
            return True
        vlm_err = str(vlm_payload)
        ok_lm, lm_payload = _try_load_lm()
        if ok_lm:
            _MODEL, _PROCESSOR = lm_payload  # type: ignore[misc]
            _CONFIG = None
            _LOAD_BACKEND = "lm"
            _LAST_LOAD_ERROR = ""
            return True
        lm_err = str(lm_payload)
    else:
        ok_lm, lm_payload = _try_load_lm()
        if ok_lm:
            _MODEL, _PROCESSOR = lm_payload  # type: ignore[misc]
            _CONFIG = None
            _LOAD_BACKEND = "lm"
            _LAST_LOAD_ERROR = ""
            return True
        lm_err = str(lm_payload)
        ok_vlm, vlm_payload = _try_load_vlm()
        if ok_vlm:
            _MODEL, _PROCESSOR, _CONFIG = vlm_payload  # type: ignore[misc]
            _LOAD_BACKEND = "vlm"
            _LAST_LOAD_ERROR = ""
            return True
        vlm_err = str(vlm_payload)

    _MODEL = None
    _PROCESSOR = None
    _CONFIG = None
    _LOAD_BACKEND = ""
    _LAST_LOAD_ERROR = f"mlx_lm load failed: {lm_err} | mlx_vlm load failed: {vlm_err}"
    return False


def _run_hf_python(code: str, timeout: int = 120) -> tuple[bool, str]:
    py = _mlx_python()
    if not os.path.isfile(py):
        return False, "Utility MLX Python not found."
    try:
        r = subprocess.run(
            [py, "-c", code],
            capture_output=True,
            text=True,
            timeout=timeout,
            env={**os.environ, "HF_HUB_ENABLE_HF_TRANSFER": "1"},
        )
        if r.returncode != 0:
            return False, (r.stderr or r.stdout or "Command failed").strip()[:500]
        return True, (r.stdout or "").strip()
    except Exception as e:
        return False, str(e)


def ensure_downloaded(on_progress: Callable[[str], None] | None = None) -> tuple[bool, str]:
    """Ensure utility MLX venv exists and utility model snapshot is available locally."""
    repo = _get_model_repo()
    if on_progress:
        on_progress("Preparing utility MLX environment...")
    ok, msg = ensure_utility_mlx_venv()
    if not ok:
        return False, msg

    local_path = _local_model_path_if_any(repo)
    if local_path is not None:
        if os.path.exists(local_path):
            return True, "Local utility model path is available."
        return False, f"Local utility model path not found: {local_path}"

    if on_progress:
        on_progress(f"Checking cached utility model ({repo})...")
    check_code = (
        "from huggingface_hub import snapshot_download\n"
        f"snapshot_download({repo!r}, local_files_only=True)\n"
        "print('ok')\n"
    )
    cached_ok, _ = _run_hf_python(check_code, timeout=60)
    if cached_ok:
        return True, "Utility model already downloaded."

    if on_progress:
        on_progress(f"Downloading utility model ({repo})...")
    dl_code = (
        "from huggingface_hub import snapshot_download\n"
        f"snapshot_download({repo!r})\n"
        "print('ok')\n"
    )
    dl_ok, dl_msg = _run_hf_python(dl_code, timeout=3600)
    if not dl_ok:
        return False, dl_msg
    return True, "Utility model downloaded."


def is_available() -> bool:
    """True when utility model is enabled and cached locally."""
    if not _enabled():
        return False
    repo = _get_model_repo()
    local_path = _local_model_path_if_any(repo)
    if local_path is not None:
        return os.path.exists(local_path)
    check_code = (
        "from huggingface_hub import snapshot_download\n"
        f"snapshot_download({repo!r}, local_files_only=True)\n"
        "print('ok')\n"
    )
    ok, _ = _run_hf_python(check_code, timeout=60)
    return ok


def is_loaded() -> bool:
    return _MODEL is not None and _PROCESSOR is not None and (_CONFIG is not None or _LOAD_BACKEND == "lm")


def unload() -> None:
    """Release in-memory model objects and stop subprocess worker."""
    global _MODEL, _PROCESSOR, _CONFIG, _LOAD_BACKEND
    with _LOAD_LOCK:
        _stop_subprocess_worker()
        _MODEL = None
        _PROCESSOR = None
        _CONFIG = None
        _LOAD_BACKEND = ""
        gc.collect()


def utility_model_status() -> dict:
    """Return status payload for preferences/setup surfaces."""
    worker_alive = _WORKER is not None and _WORKER.alive
    return {
        "available": is_available(),
        "loaded": is_loaded(),
        "model_repo": _get_model_repo(),
        "enabled": _enabled(),
        "utility_venv_installed": utility_mlx_venv_exists(),
        "prism_mlx_ready": utility_mlx_prism_ready(),
        "backend": _LOAD_BACKEND,
        "subprocess_worker_alive": worker_alive,
        "last_error": _LAST_LOAD_ERROR,
    }


def utility_complete(
    messages: list[dict],
    max_tokens: int = 512,
    temperature: float = 0.3,
    on_log: Callable[[str], None] | None = None,
    activity_operation: str | None = None,
    activity_thread_id: str | None = None,
    activity_job_id: str | None = None,
    surface_errors: bool = True,
) -> str | None:
    """Run a best-effort completion with the utility model.

    Returns None when disabled/unavailable/failing so callers can fallback.

    When ``surface_errors`` is False, empty/unavailable outcomes do not emit
    Activity errors — callers with their own fallback chain (e.g. knowledge
    digest) surface failures only after all paths for a source fail.
    """
    if not _enabled():
        return None
    prompt = _messages_to_prompt(messages)
    if not prompt:
        return None
    op_name, op_def = _activity_definition(activity_operation)
    job_id = (
        str(activity_job_id or "").strip()
        or f"{op_name}:{int(time.time() * 1000)}:{threading.get_ident()}"
    )
    quiet_activity = not surface_errors
    if not quiet_activity:
        _emit_activity(
            "started",
            op_def["started"],
            operation=op_name,
            job_id=job_id,
            thread_id=activity_thread_id,
        )
    with _LOAD_LOCK:
        if not _ensure_loaded():
            # In-process model load failed (likely ABI mismatch between
            # backend Python and utility venv Python).  Fall back to the
            # persistent subprocess worker that runs the venv's own interpreter.
            text = _subprocess_utility_complete(messages, max_tokens, temperature)
            if text:
                _emit_activity(
                    "completed",
                    f"{op_def['completed']} (subprocess bridge)",
                    operation=op_name,
                    job_id=job_id,
                    thread_id=activity_thread_id,
                )
                return text
            if surface_errors:
                _emit_activity(
                    "error",
                    f"{op_def['error']} (model unavailable; subprocess bridge also failed)",
                    operation=op_name,
                    job_id=job_id,
                    thread_id=activity_thread_id,
                )
            return None
        try:
            t0 = time.monotonic()
            text = ""
            for attempt, attempt_messages in enumerate((messages, _utility_retry_messages(messages))):
                raw = _generate_inprocess_raw(
                    attempt_messages,
                    prompt=_messages_to_prompt(attempt_messages),
                    max_tokens=max_tokens,
                    temperature=temperature,
                )
                text = _normalize_utility_output(raw)
                if text or attempt == 1:
                    break
            elapsed = time.monotonic() - t0
            if on_log:
                on_log(f"utility_complete finished in {elapsed:.2f}s")
            if text:
                _emit_activity(
                    "completed",
                    f"{op_def['completed']} ({elapsed:.1f}s)",
                    operation=op_name,
                    job_id=job_id,
                    thread_id=activity_thread_id,
                )
                return text
            if surface_errors:
                _emit_activity(
                    "error",
                    f"{op_def['error']} (empty result)",
                    operation=op_name,
                    job_id=job_id,
                    thread_id=activity_thread_id,
                )
            return None
        except Exception as exc:
            if surface_errors:
                err = str(exc).strip()
                suffix = f": {err[:120]}" if err else ""
                _emit_activity(
                    "error",
                    f"{op_def['error']}{suffix}",
                    operation=op_name,
                    job_id=job_id,
                    thread_id=activity_thread_id,
                )
            return None


def _generate_inprocess_raw(
    messages: list[dict],
    *,
    prompt: str,
    max_tokens: int,
    temperature: float,
) -> str:
    if _LOAD_BACKEND == "lm":
        from mlx_lm import generate as generate_lm

        effective_prompt = _apply_utility_chat_template(_PROCESSOR, messages) or prompt
        try:
            result = generate_lm(
                _MODEL,
                _PROCESSOR,
                prompt=effective_prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                verbose=False,
            )
        except TypeError:
            result = generate_lm(
                _MODEL,
                _PROCESSOR,
                prompt=effective_prompt,
                max_tokens=max_tokens,
                verbose=False,
            )
    else:
        from mlx_vlm import generate
        from mlx_vlm.prompt_utils import apply_chat_template

        try:
            formatted_prompt = apply_chat_template(
                _PROCESSOR,
                _CONFIG,
                prompt,
                num_images=0,
                enable_thinking=False,
            )
        except TypeError:
            formatted_prompt = apply_chat_template(_PROCESSOR, _CONFIG, prompt, num_images=0)

        result = generate(
            _MODEL,
            _PROCESSOR,
            formatted_prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            verbose=False,
        )
    return _decode_generate_result(result)


def _hash_embed_text(text: str, *, dimensions: int) -> list[float]:
    """Deterministic lightweight embedding for utility-model paths.

    This keeps session-search embeddings inside the utility-model subsystem
    without introducing an additional model dependency or venv.
    """
    dims = max(8, min(int(dimensions), 2048))
    vec = [0.0] * dims
    tokens = re.findall(r"[A-Za-z0-9_]+", (text or "").lower())
    if not tokens:
        return vec

    def _accum_feature(feature: str, weight: float) -> None:
        h = hashlib.sha256(feature.encode("utf-8", errors="ignore")).digest()
        idx = int.from_bytes(h[:4], "big", signed=False) % dims
        sign = -1.0 if (h[4] & 1) else 1.0
        vec[idx] += sign * weight

    for i, tok in enumerate(tokens):
        # token feature
        _accum_feature(f"tok:{tok}", 1.0)
        # adjacent bigram feature to preserve local phrase meaning
        if i + 1 < len(tokens):
            _accum_feature(f"bi:{tok}_{tokens[i + 1]}", 0.7)
        # subword trigrams improve wording-mismatch robustness
        padded = f"#{tok}#"
        for j in range(max(0, len(padded) - 2)):
            _accum_feature(f"tri:{padded[j:j + 3]}", 0.35)
    norm = math.sqrt(sum(v * v for v in vec))
    if norm > 0:
        inv = 1.0 / norm
        vec = [v * inv for v in vec]
    return vec


def utility_embedding_identity() -> str:
    """Active embedding identity (mlx:repo:dim or hash:N)."""
    try:
        from utils.embedding_model import embedding_identity

        return embedding_identity()
    except Exception:
        return "hash:256"


def utility_embed_texts(
    texts: list[str],
    *,
    dimensions: int = 256,
    kind: str = "document",
    activity_operation: str | None = "session_search_embedding",
    activity_thread_id: str | None = None,
    activity_job_id: str | None = None,
    emit_activity: bool = False,
) -> list[list[float]] | None:
    """Best-effort local embeddings: MLX embedder first, hash fallback when unavailable.

    Returns None when both paths fail so callers can degrade to lexical-only retrieval.
    """
    if not isinstance(texts, list):
        return None
    cleaned = [str(t or "").strip() for t in texts]
    if not cleaned:
        return []
    embed_kind = "query" if str(kind or "").strip().lower() == "query" else "document"
    try:
        from utils.embedding_model import embed_texts as mlx_embed_texts

        mlx_out = mlx_embed_texts(cleaned, kind=embed_kind)
        if mlx_out is not None and len(mlx_out) == len(cleaned):
            return mlx_out
    except Exception:
        pass
    # Hash fallback when MLX is not loaded. If the model is loaded but forward failed,
    # do not emit hash vectors at MLX dimension under an MLX identity stamp.
    try:
        from utils.embedding_model import _hash_fallback_dim, embedding_enabled, is_real

        if not embedding_enabled():
            return None
        if is_real():
            return None
        hash_dim = int(_hash_fallback_dim())
    except Exception:
        hash_dim = int(dimensions)
    op_name, op_def = _activity_definition(activity_operation)
    job_id = (
        str(activity_job_id or "").strip()
        or f"{op_name}:{int(time.time() * 1000)}:{threading.get_ident()}"
    )
    if emit_activity and _enabled():
        _emit_activity(
            "started",
            op_def["started"],
            operation=op_name,
            job_id=job_id,
            thread_id=activity_thread_id,
        )
    try:
        out = [_hash_embed_text(t, dimensions=hash_dim) for t in cleaned]
        if emit_activity and _enabled():
            _emit_activity(
                "completed",
                op_def["completed"],
                operation=op_name,
                job_id=job_id,
                thread_id=activity_thread_id,
            )
        return out
    except Exception as exc:
        err = str(exc).strip()
        suffix = f": {err[:120]}" if err else ""
        if emit_activity and _enabled():
            _emit_activity(
                "error",
                f"{op_def['error']}{suffix}",
                operation=op_name,
                job_id=job_id,
                thread_id=activity_thread_id,
            )
        return None
