"""MLX embedding model for knowledge/session search (hash fallback when unavailable)."""
from __future__ import annotations

import logging
import math
import os
import site
import subprocess
import sys
import threading
import time
from typing import Any, Literal

_log = logging.getLogger(__name__)

EmbedKind = Literal["document", "query"]

# mlx-community/bge-small-en-v1.5-mlx was never published; bf16/4bit are the real MLX ports.
_DEFAULT_REPO = "mlx-community/bge-small-en-v1.5-bf16"
_LEGACY_EMBEDDING_REPOS: dict[str, str] = {
    "mlx-community/bge-small-en-v1.5-mlx": _DEFAULT_REPO,
}
_HASH_IDENTITY_DIM = 256
_LOAD_BACKOFF_SEC = 300.0
_MAX_SEQ_LEN = 512
_BATCH_MAX = 16
_SUBWINDOW_CHARS = 1800
_IDLE_UNLOAD_SEC = 600.0

_LOAD_LOCK = threading.Lock()
_MODEL: Any = None
_TOKENIZER: Any = None
_REPO: str = ""
_DIM: int = 0
_POOLING: str = "mean"
_FAMILY: str = "generic"
_LOAD_FAILED_UNTIL: float = 0.0
_LAST_USED: float = 0.0
_LAST_LOAD_ERROR: str = ""
_MLX_EMBED_PKGS_READY: bool = False
_ACTIVITY_ERR_DEDUPE_SEC = 45.0
_last_reindex_activity_err: tuple[str, float] = ("", 0.0)

_reindex_lock = threading.Lock()
_reindex_inflight: set[str] = set()
_reindex_status: dict[str, Any] = {"state": "idle", "corpora_pending": 0, "current_corpus": ""}
_reindex_hold_lock = threading.Lock()
_reindex_hold_count = 0
_reindex_yield_lock = threading.Lock()
_reindex_yield_requested = False
_startup_sweep_done = False
_graph_only_failures: dict[str, int] = {}
_MAX_GRAPH_ONLY_RETRIES = 8


def _utility_cfg() -> dict[str, Any]:
    try:
        from core.config import get_utility_config

        return get_utility_config() or {}
    except Exception:
        return {}


def embedding_enabled() -> bool:
    cfg = _utility_cfg()
    return bool(cfg.get("embedding_enabled", True))


def _configured_repo() -> str:
    cfg = _utility_cfg()
    repo = str(cfg.get("embedding_model_repo") or "").strip() or _DEFAULT_REPO
    return _LEGACY_EMBEDDING_REPOS.get(repo, repo)


def _local_embedding_model_path(model_ref: str) -> str | None:
    """Expanded local path when *model_ref* points at the 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 _resolve_embedding_repo() -> tuple[str, str | None]:
    """Return ``(repo_for_mlx_load, error_message)``."""
    repo = _configured_repo()
    local = _local_embedding_model_path(repo)
    if local is not None:
        if os.path.isdir(local) or os.path.isfile(local):
            return local, None
        return repo, (
            f"embedding model path not found: {local} "
            "(fix [utility] embedding_model_repo in Settings or mugi_config.toml)"
        )
    return repo, None


def _mlx_venv_site_packages() -> list[str]:
    """Site-packages dirs for the main MLX venv (where mlx-embeddings is installed)."""
    try:
        from app_mlx import MLX_VENV_DIR
    except Exception:
        return []
    paths: list[str] = []
    lib_root = os.path.join(MLX_VENV_DIR, "lib")
    try:
        for entry in sorted(os.listdir(lib_root)):
            if entry.startswith("python"):
                paths.append(os.path.join(lib_root, entry, "site-packages"))
    except OSError:
        pass
    ver = f"python{sys.version_info.major}.{sys.version_info.minor}"
    paths.append(os.path.join(MLX_VENV_DIR, "lib", ver, "site-packages"))
    try:
        paths.extend(site.getsitepackages([MLX_VENV_DIR]))
    except Exception:
        pass
    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 _mlx_venv_python() -> str:
    try:
        from app_mlx import MLX_VENV_DIR

        py = os.path.join(MLX_VENV_DIR, "bin", "python")
        if os.path.isfile(py):
            return py
    except Exception:
        pass
    return ""


def _run_mlx_venv_python(code: str, *, timeout: int = 3600) -> tuple[bool, str]:
    py = _mlx_venv_python()
    if not py:
        return False, "MLX venv Python not found (~/.mugi/mlx-venv)."
    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_embedding_repo_cached(repo: str) -> tuple[bool, str]:
    """Download the MLX embedding weights into the Hugging Face cache if missing."""
    ref = (repo or "").strip()
    if not ref or ref.startswith(("/", "~")):
        return True, ""
    check = (
        "from huggingface_hub import snapshot_download\n"
        f"snapshot_download({ref!r}, local_files_only=True)\n"
        "print('ok')\n"
    )
    ok, _ = _run_mlx_venv_python(check, timeout=120)
    if ok:
        return True, ""
    dl = (
        "from huggingface_hub import snapshot_download\n"
        f"snapshot_download({ref!r})\n"
        "print('ok')\n"
    )
    return _run_mlx_venv_python(dl, timeout=3600)


def _is_model_not_found_error(exc: BaseException) -> bool:
    name = type(exc).__name__
    if "ModelNotFound" in name:
        return True
    low = str(exc).lower()
    return "model not found" in low or "repository not found" in low


def _ensure_mlx_embeddings_importable() -> bool:
    """Add MLX venv site-packages so backend can import mlx-embeddings from the MLX stack."""
    global _MLX_EMBED_PKGS_READY
    if _MLX_EMBED_PKGS_READY:
        return True
    for sp in _mlx_venv_site_packages():
        if sp not in sys.path:
            sys.path.insert(0, sp)
    try:
        import mlx_embeddings  # noqa: F401

        _MLX_EMBED_PKGS_READY = True
        return True
    except ImportError:
        return False


def embedding_load_error() -> str:
    """Last load failure reason (empty when load succeeded or not yet attempted)."""
    return str(_LAST_LOAD_ERROR or "").strip()


def _hash_fallback_dim() -> int:
    try:
        from core.config import get_session_search_config

        raw = get_session_search_config().get("semantic_embedding_dimensions", _HASH_IDENTITY_DIM)
        return max(32, min(int(raw), 2048))
    except Exception:
        return _HASH_IDENTITY_DIM


def _detect_family(repo: str) -> tuple[str, str]:
    low = (repo or "").lower()
    if "bge" in low or "en-v1.5" in low:
        return "bge", "cls"
    if "nomic" in low:
        return "nomic", "mean"
    if "e5" in low or "e5-" in low:
        return "e5", "mean"
    if "gte" in low:
        return "gte", "mean"
    if "minilm" in low:
        return "minilm", "mean"
    return "generic", "mean"


def _prefix_for(kind: EmbedKind, family: str, cfg: dict[str, Any]) -> str:
    q_override = str(cfg.get("embedding_query_prefix") or "").strip()
    d_override = str(cfg.get("embedding_doc_prefix") or "").strip()
    if kind == "query" and q_override:
        return q_override
    if kind == "document" and d_override:
        return d_override
    if family == "bge":
        if kind == "query":
            return "Represent this sentence for searching relevant passages: "
        return ""
    if family == "nomic":
        return "search_query: " if kind == "query" else "search_document: "
    if family == "e5":
        return "query: " if kind == "query" else "passage: "
    return ""


def _apply_prefix(text: str, *, kind: EmbedKind) -> str:
    cfg = _utility_cfg()
    prefix = _prefix_for(kind, _FAMILY, cfg)
    body = (text or "").strip()
    if not body:
        return ""
    return f"{prefix}{body}" if prefix else body


def _l2_normalize(vec: list[float]) -> list[float]:
    norm = math.sqrt(sum(v * v for v in vec))
    if norm <= 0:
        return vec
    return [v / norm for v in vec]


def _mean_pool(hidden: Any, mask: Any | None) -> Any:
    import mlx.core as mx

    if mask is None:
        return mx.mean(hidden, axis=1)
    mask_f = mask.astype(mx.float32)
    summed = mx.sum(hidden * mask_f[:, :, None], axis=1)
    counts = mx.maximum(mx.sum(mask_f, axis=1, keepdims=True), 1e-6)
    return summed / counts


def _pool_hidden(hidden: Any, mask: Any | None, pooling: str) -> Any:
    import mlx.core as mx

    if pooling == "cls":
        return hidden[:, 0, :]
    return _mean_pool(hidden, mask)


def _hf_tokenizer():
    """mlx-embeddings returns a TokenizerWrapper that is not callable; use the inner HF tokenizer."""
    tok = _TOKENIZER
    if tok is None:
        return None
    inner = getattr(tok, "_tokenizer", None)
    return inner if inner is not None else tok


def _forward_batch(texts: list[str]) -> list[list[float]] | None:
    global _MODEL, _TOKENIZER
    if _MODEL is None or _TOKENIZER is None:
        return None
    tokenizer = _hf_tokenizer()
    if tokenizer is None:
        return None
    try:
        import mlx.core as mx
    except ImportError:
        return None
    try:
        inputs = tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=_MAX_SEQ_LEN,
            return_tensors="np",
        )
        input_ids = mx.array(inputs["input_ids"])
        attention_mask = mx.array(inputs["attention_mask"]) if "attention_mask" in inputs else None
        try:
            outputs = _MODEL(input_ids, attention_mask=attention_mask)
        except TypeError:
            outputs = _MODEL(input_ids)
        hidden = getattr(outputs, "last_hidden_state", None)
        if hidden is None and isinstance(outputs, (tuple, list)) and outputs:
            hidden = outputs[0]
        if hidden is None:
            text_embeds = getattr(outputs, "text_embeds", None)
            if text_embeds is not None:
                pooled = text_embeds
            else:
                return None
        else:
            pooled = _pool_hidden(hidden, attention_mask, _POOLING)
        if hasattr(pooled, "tolist"):
            rows = pooled.tolist()
        else:
            rows = [list(r) for r in pooled]
        return [_l2_normalize([float(x) for x in row]) for row in rows]
    except Exception as exc:
        _log.warning("[embedding] forward failed: %s", str(exc)[:120])
        return None


def _subwindows(text: str) -> list[str]:
    t = (text or "").strip()
    if not t:
        return []
    if len(t) <= _SUBWINDOW_CHARS:
        return [t]
    out: list[str] = []
    start = 0
    overlap = 200
    step = max(1, _SUBWINDOW_CHARS - overlap)
    while start < len(t):
        chunk = t[start : start + _SUBWINDOW_CHARS].strip()
        if chunk:
            out.append(chunk)
        start += step
    return out or [t[:_SUBWINDOW_CHARS]]


def _average_vectors(vecs: list[list[float]]) -> list[float]:
    if not vecs:
        return []
    if len(vecs) == 1:
        return list(vecs[0])
    dim = len(vecs[0])
    acc = [0.0] * dim
    for v in vecs:
        n = min(dim, len(v))
        for i in range(n):
            acc[i] += float(v[i])
    inv = 1.0 / len(vecs)
    return _l2_normalize([x * inv for x in acc])


def load() -> bool:
    """Lazy-load the MLX embedding model (independent of utility generative model)."""
    global _MODEL, _TOKENIZER, _REPO, _DIM, _POOLING, _FAMILY, _LOAD_FAILED_UNTIL, _LAST_USED, _LAST_LOAD_ERROR
    if not embedding_enabled():
        _LAST_LOAD_ERROR = ""
        return False
    now = time.time()
    if _MODEL is not None and _TOKENIZER is not None:
        _LAST_USED = now
        _LAST_LOAD_ERROR = ""
        return True
    if now < _LOAD_FAILED_UNTIL:
        return False
    with _LOAD_LOCK:
        if _MODEL is not None and _TOKENIZER is not None:
            _LAST_USED = now
            _LAST_LOAD_ERROR = ""
            return True
        if now < _LOAD_FAILED_UNTIL:
            return False
        repo, resolve_err = _resolve_embedding_repo()
        if resolve_err:
            _LAST_LOAD_ERROR = resolve_err
            _log.warning("[embedding] %s", resolve_err)
            _LOAD_FAILED_UNTIL = time.time() + _LOAD_BACKOFF_SEC
            return False
        if not _ensure_mlx_embeddings_importable():
            _LAST_LOAD_ERROR = (
                "mlx-embeddings is not installed in ~/.mugi/mlx-venv. "
                "Settings → LLMs → Revert to stable or Reinstall environment (pinned stack includes mlx-embeddings)."
            )
            _log.warning("[embedding] %s", _LAST_LOAD_ERROR)
            _LOAD_FAILED_UNTIL = time.time() + _LOAD_BACKOFF_SEC
            return False
        try:
            from mlx_embeddings.utils import load as mlx_load
        except ImportError:
            try:
                from mlx_embeddings import load as mlx_load
            except ImportError:
                _LAST_LOAD_ERROR = (
                    "mlx-embeddings import failed after MLX venv path setup. "
                    "Re-run Settings → LLMs → Update MLX stack."
                )
                _log.warning("[embedding] %s", _LAST_LOAD_ERROR)
                _LOAD_FAILED_UNTIL = time.time() + _LOAD_BACKOFF_SEC
                return False
        def _try_mlx_load() -> tuple[Any, Any]:
            try:
                return mlx_load(repo, trust_remote_code=False)
            except TypeError:
                return mlx_load(repo)

        try:
            model, tokenizer = _try_mlx_load()
        except Exception as exc:
            if _is_model_not_found_error(exc):
                ok_dl, dl_msg = _ensure_embedding_repo_cached(repo)
                if ok_dl:
                    try:
                        model, tokenizer = _try_mlx_load()
                    except Exception as retry_exc:
                        exc = retry_exc
                    else:
                        exc = None  # type: ignore[assignment]
                else:
                    _LAST_LOAD_ERROR = (
                        f"embedding model not cached ({repo}): {dl_msg[:160]}. "
                        "Log in to Hugging Face if needed (huggingface-cli login), then retry."
                    )
                    _log.warning("[embedding] %s", _LAST_LOAD_ERROR)
                    _LOAD_FAILED_UNTIL = time.time() + _LOAD_BACKOFF_SEC
                    return False
            if exc is not None:
                _LAST_LOAD_ERROR = f"embedding model load failed: {str(exc)[:200]}"
                _log.warning("[embedding] load failed for %s: %s", repo, str(exc)[:160])
                _LOAD_FAILED_UNTIL = time.time() + _LOAD_BACKOFF_SEC
                return False
        family, pooling = _detect_family(repo)
        dim = 0
        try:
            if hasattr(model, "config") and getattr(model.config, "hidden_size", None):
                dim = int(model.config.hidden_size)
        except Exception:
            dim = 0
        if dim <= 0:
            probe = _forward_batch_after_load(model, tokenizer, pooling)
            if probe:
                dim = len(probe[0])
        if dim <= 0:
            dim = 384
        _MODEL = model
        _TOKENIZER = tokenizer
        _REPO = repo
        _DIM = dim
        _POOLING = pooling
        _FAMILY = family
        _LOAD_FAILED_UNTIL = 0.0
        _LAST_USED = time.time()
        _LAST_LOAD_ERROR = ""
        _log.info("[embedding] loaded %s dim=%d pooling=%s", repo, dim, pooling)
        return True


def _forward_batch_after_load(model: Any, tokenizer: Any, pooling: str) -> list[list[float]] | None:
    global _MODEL, _TOKENIZER, _POOLING
    _MODEL = model
    _TOKENIZER = tokenizer
    _POOLING = pooling
    return _forward_batch(["embedding dimension probe"])


def _maybe_idle_unload() -> None:
    global _MODEL, _TOKENIZER
    if _MODEL is None:
        return
    if time.time() - _LAST_USED < _IDLE_UNLOAD_SEC:
        return
    with _LOAD_LOCK:
        if _MODEL is None or time.time() - _LAST_USED < _IDLE_UNLOAD_SEC:
            return
        _MODEL = None
        _TOKENIZER = None
        _log.debug("[embedding] idle-unloaded model")


def _foreground_chat_busy() -> bool:
    try:
        from persistence import read_main_chat_busy

        return read_main_chat_busy()
    except Exception:
        return False


def embedding_work_should_defer() -> bool:
    """True when background embedding/index writes must yield to foreground chat."""
    if _foreground_chat_busy():
        return True
    with _reindex_yield_lock:
        requested = bool(_reindex_yield_requested)
    if not requested:
        return False
    # Yield is only meaningful while a reindex hold is still active; once the
    # runner exits, a stale flag must not block later background recovery.
    if not _embedding_reindex_in_progress():
        _clear_embedding_yield_to_chat()
        return False
    return True


def _request_embedding_yield_to_chat() -> None:
    with _reindex_yield_lock:
        global _reindex_yield_requested
        _reindex_yield_requested = True


def _clear_embedding_yield_to_chat() -> None:
    with _reindex_yield_lock:
        global _reindex_yield_requested
        _reindex_yield_requested = False


def begin_embedding_reindex_hold() -> None:
    """Keep MLX embedder loaded while a corpus reindex is in flight."""
    global _reindex_hold_count
    with _reindex_hold_lock:
        _reindex_hold_count += 1


def end_embedding_reindex_hold() -> None:
    """Release reindex hold started by :func:`begin_embedding_reindex_hold`."""
    global _reindex_hold_count
    with _reindex_hold_lock:
        _reindex_hold_count = max(0, _reindex_hold_count - 1)


def _embedding_reindex_in_progress() -> bool:
    with _reindex_hold_lock:
        return _reindex_hold_count > 0


def release_for_foreground_chat() -> None:
    """Unload in-process MLX embedder so foreground chat can use GPU (best-effort)."""
    global _MODEL, _TOKENIZER
    if _embedding_reindex_in_progress():
        _request_embedding_yield_to_chat()
        _log.debug("[embedding] yield requested — background reindex defers to chat")
    with _LOAD_LOCK:
        if _MODEL is None:
            return
        _MODEL = None
        _TOKENIZER = None
        _log.debug("[embedding] released for foreground chat")


def is_real() -> bool:
    return _MODEL is not None and _TOKENIZER is not None and _DIM > 0


def embedding_dim() -> int:
    if is_real():
        return int(_DIM)
    return _hash_fallback_dim()


def embedding_identity() -> str:
    if is_real():
        return f"mlx:{_REPO}:{_DIM}"
    return f"hash:{_hash_fallback_dim()}"


def _configured_embedding_dim() -> int:
    try:
        from core.config import get_session_search_config

        raw = get_session_search_config().get("semantic_embedding_dimensions")
        if raw is not None:
            return max(32, min(int(raw), 2048))
    except Exception:
        pass
    return 384


def configured_embedding_identity() -> str:
    """Expected MLX identity from settings (stable while the model is idle-unloaded)."""
    if not embedding_enabled():
        return f"hash:{_hash_fallback_dim()}"
    return f"mlx:{_configured_repo()}:{_configured_embedding_dim()}"


def _parse_embedding_identity(identity: str) -> tuple[str, str, int]:
    text = (identity or "").strip()
    if not text:
        return "", "", 0
    if text.startswith("hash:"):
        try:
            return "hash", "", int(text.split(":", 1)[1])
        except ValueError:
            return "hash", "", _hash_fallback_dim()
    if text.startswith("mlx:"):
        body = text[4:]
        repo, _, dim_s = body.rpartition(":")
        try:
            return "mlx", repo, int(dim_s)
        except ValueError:
            return "mlx", body, 0
    return "", text, 0


def _same_configured_mlx_repo(stored_repo: str, expected_repo: str) -> bool:
    if not stored_repo or not expected_repo:
        return False
    if stored_repo == expected_repo:
        return True
    return _LEGACY_EMBEDDING_REPOS.get(stored_repo, stored_repo) == expected_repo


def embedding_identity_stale(stored: str) -> bool:
    """True when a stamped identity no longer matches configured or loaded MLX settings."""
    stamp = (stored or "").strip()
    if not stamp:
        return False
    if is_real():
        return stamp != embedding_identity()
    if not embedding_enabled():
        return not stamp.startswith("hash:")
    expected = configured_embedding_identity()
    if stamp == expected:
        return False
    kind, repo, _dim = _parse_embedding_identity(stamp)
    exp_kind, exp_repo, _exp_dim = _parse_embedding_identity(expected)
    if kind == "mlx" and exp_kind == "mlx":
        return not _same_configured_mlx_repo(repo, exp_repo)
    return True


def embed_texts(
    texts: list[str],
    *,
    kind: EmbedKind = "document",
    normalize: bool = True,
) -> list[list[float]] | None:
    """Embed texts with the MLX model. Returns None on failure (caller uses hash fallback)."""
    if not texts:
        return []
    if not embedding_enabled():
        return None
    if not load():
        return None
    prefixed = [_apply_prefix(t, kind=kind) for t in texts]
    out: list[list[float]] = []
    for raw, pref in zip(texts, prefixed):
        if not pref.strip():
            out.append([])
            continue
        windows = [_apply_prefix(w, kind=kind) for w in _subwindows(raw)]
        window_vecs: list[list[float]] = []
        for i in range(0, len(windows), _BATCH_MAX):
            batch = [w for w in windows[i : i + _BATCH_MAX] if w.strip()]
            if not batch:
                continue
            batch_vecs = _forward_batch(batch)
            if not batch_vecs:
                _maybe_idle_unload()
                return None
            window_vecs.extend(batch_vecs)
        merged = _average_vectors(window_vecs)
        if not merged:
            out.append([])
            continue
        out.append(_l2_normalize(merged) if normalize else merged)
    _maybe_idle_unload()
    if len(out) != len(texts):
        return None
    return out


def get_embedding_reindex_status() -> dict[str, Any]:
    with _reindex_lock:
        out = dict(_reindex_status)
    err = embedding_load_error()
    if err:
        out["load_error"] = err
    return out


def _set_reindex_status(**kwargs: object) -> None:
    with _reindex_lock:
        _reindex_status.update(kwargs)


def _corpus_slug_from_thread(thread_id: str) -> str | None:
    tid = (thread_id or "").strip()
    if not tid.startswith("knowledge:"):
        return None
    slug = tid.split(":", 1)[1].strip()
    return slug or None


def active_reindex_staging_thread_ids() -> set[str]:
    """Staging thread ids for corpora with a reindex worker in flight."""
    try:
        import session_search_db as ss
    except Exception:
        return set()
    with _reindex_lock:
        return {
            ss.reindex_staging_thread_id(tid)
            for tid in _reindex_inflight
            if str(tid).startswith("knowledge:")
        }


def _search_index_current_for_corpus(thread_id: str) -> bool:
    """True when production search vectors + identity are up to date (graph-only retry)."""
    try:
        import session_search_db as ss
        from utils.embedding_chunking import thread_chunker_version_stale

        conn = ss._safe_get_connection()
        if conn is None:
            return False
        try:
            if not ss.knowledge_semantic_index_complete(conn, thread_id):
                return False
            if thread_chunker_version_stale(conn, thread_id):
                return False
            stored = ss._get_meta(conn, ss._embed_identity_meta_key(thread_id))
            if not stored:
                return False
            return not embedding_identity_stale(stored)
        finally:
            conn.close()
    except Exception as exc:
        _log.debug(
            "[embedding] search index current check failed for %s: %s",
            thread_id,
            str(exc)[:120],
        )
        return False


def thread_needs_embedding_reindex(thread_id: str) -> bool:
    """True when stored identity or global chunker_version is stale for a knowledge corpus."""
    slug = _corpus_slug_from_thread(thread_id)
    if not slug:
        return False
    try:
        import session_search_db as ss
        from utils.embedding_chunking import thread_chunker_version_stale

        conn = ss._safe_get_connection()
        if conn is None:
            return False
        try:
            fts_row = conn.execute(
                "SELECT COUNT(*) FROM messages_fts WHERE thread_id = ?",
                (thread_id,),
            ).fetchone()
            fts_n = int(fts_row[0] or 0) if fts_row else 0
            if fts_n <= 0:
                return False
            if thread_chunker_version_stale(conn, thread_id):
                return True
            if not ss.knowledge_semantic_index_complete(conn, thread_id):
                return True
            stored = ss._get_meta(conn, ss._embed_identity_meta_key(thread_id))
            if not stored:
                # Complete vectors without a stamp — do not trigger destructive rebuild.
                return False
            return embedding_identity_stale(stored)
        finally:
            conn.close()
    except Exception as exc:
        _log.debug("[embedding] thread_needs_reindex check failed: %s", str(exc)[:120])
        return False


def maybe_schedule_embedding_reindex(thread_id: str) -> bool:
    """Enqueue a background reindex for one knowledge corpus thread (single-flight)."""
    slug = _corpus_slug_from_thread(thread_id)
    if not slug:
        return False
    if _foreground_chat_busy():
        _log.debug("[embedding] reindex deferred — main chat active (%s)", slug)

        def _retry() -> None:
            maybe_schedule_embedding_reindex(thread_id)

        threading.Timer(90.0, _retry).start()
        return False
    with _reindex_lock:
        if thread_id in _reindex_inflight or slug in _reindex_inflight:
            return False
        _reindex_inflight.add(thread_id)
        _reindex_inflight.add(slug)
        pending = int(_reindex_status.get("corpora_pending") or 0) + 1
        _reindex_status.update(
            {"state": "reindexing", "corpora_pending": pending, "current_corpus": slug}
        )

    # Cross-process single-owner for this corpus's reindex: ``_reindex_inflight`` only
    # dedupes within one process, but the backend, a stale runner, MCP, and CLI all reindex
    # against the same ``~/.mugi/search.db`` and ``knowledge:{slug}`` staging thread.
    from utils.process_lock import CrossProcessLock

    _xproc_lock = CrossProcessLock(f"embed_reindex_{slug}")
    if not _xproc_lock.try_acquire():
        _log.debug("[embedding] reindex skipped — another process owns slug %s", slug)
        with _reindex_lock:
            _reindex_inflight.discard(thread_id)
            _reindex_inflight.discard(slug)
            pending_now = max(0, int(_reindex_status.get("corpora_pending") or 1) - 1)
            if pending_now <= 0 and not _reindex_inflight:
                _reindex_status.update(
                    {"state": "idle", "corpora_pending": 0, "current_corpus": ""}
                )
            else:
                _reindex_status["corpora_pending"] = pending_now
        return False

    def _runner() -> None:
        if embedding_work_should_defer():
            _log.debug("[embedding] reindex runner deferred — chat active (%s)", slug)

            def _retry_runner() -> None:
                maybe_schedule_embedding_reindex(thread_id)

            threading.Timer(90.0, _retry_runner).start()
            with _reindex_lock:
                _reindex_inflight.discard(thread_id)
                _reindex_inflight.discard(slug)
                pending_now = max(0, int(_reindex_status.get("corpora_pending") or 1) - 1)
                if pending_now <= 0 and not _reindex_inflight:
                    _reindex_status.update(
                        {"state": "idle", "corpora_pending": 0, "current_corpus": ""}
                    )
                else:
                    _reindex_status["corpora_pending"] = pending_now
            _xproc_lock.release()
            return
        _clear_embedding_yield_to_chat()
        begin_embedding_reindex_hold()
        result: dict[str, Any] = {"ok": False}
        try:
            result = reindex_embeddings_for_corpus(slug, emit_activity=True)
        finally:
            end_embedding_reindex_hold()
            _clear_embedding_yield_to_chat()
            with _reindex_lock:
                _reindex_inflight.discard(thread_id)
                _reindex_inflight.discard(slug)
                pending_now = max(0, int(_reindex_status.get("corpora_pending") or 1) - 1)
                if pending_now <= 0 and not _reindex_inflight:
                    _reindex_status.update(
                        {"state": "idle", "corpora_pending": 0, "current_corpus": ""}
                    )
                else:
                    _reindex_status["corpora_pending"] = pending_now
            _xproc_lock.release()
        if not result.get("ok"):
            err = str(result.get("error") or "")
            if err == "no_indexable_chunks":
                needs_retry = False
            elif err == "chat_busy_deferred":
                needs_retry = True
            else:
                needs_retry = thread_needs_embedding_reindex(thread_id)
                if not needs_retry and result.get("search_index_ok"):
                    with _reindex_lock:
                        fails = _graph_only_failures.get(thread_id, 0) + 1
                        _graph_only_failures[thread_id] = fails
                    if fails >= _MAX_GRAPH_ONLY_RETRIES:
                        _emit_reindex_activity_error(
                            emit_activity=True,
                            job_id=f"embed-reindex:{slug}:graph-cap",
                            thread_id=thread_id,
                            err=str(result.get("error") or "graph_embedding_failed"),
                        )
                        needs_retry = False
                    else:
                        needs_retry = True
            if needs_retry:
                delay_s = 120.0 if _foreground_chat_busy() else 90.0

                def _retry_failed() -> None:
                    maybe_schedule_embedding_reindex(thread_id)

                threading.Timer(delay_s, _retry_failed).start()
        else:
            with _reindex_lock:
                _graph_only_failures.pop(thread_id, None)

    try:
        threading.Thread(target=_runner, daemon=True, name=f"embed-reindex-{slug}").start()
    except Exception:
        # Thread creation failed (resource exhaustion): _runner will never run, so release
        # the cross-process lock and clear the in-process inflight markers here, otherwise
        # this slug's reindex is wedged until process death.
        _log.warning("[embedding] reindex thread failed to start for %s", slug, exc_info=True)
        _xproc_lock.release()
        with _reindex_lock:
            _reindex_inflight.discard(thread_id)
            _reindex_inflight.discard(slug)
            pending_now = max(0, int(_reindex_status.get("corpora_pending") or 1) - 1)
            if pending_now <= 0 and not _reindex_inflight:
                _reindex_status.update(
                    {"state": "idle", "corpora_pending": 0, "current_corpus": ""}
                )
            else:
                _reindex_status["corpora_pending"] = pending_now
        return False
    return True


def request_embedding_reindex(slug: str) -> dict[str, Any]:
    """HTTP/MCP entry: enqueue single-flight background reindex for one corpus."""
    slug = (slug or "").strip()
    if not slug:
        return {"ok": False, "error": "missing_slug", "scheduled": False}
    thread_id = f"knowledge:{slug}"
    with _reindex_lock:
        already = thread_id in _reindex_inflight or slug in _reindex_inflight
    if already:
        return {
            "ok": True,
            "slug": slug,
            "scheduled": False,
            "reason": "already_inflight",
            "embedding_reindex": get_embedding_reindex_status(),
        }
    scheduled = maybe_schedule_embedding_reindex(thread_id)
    if scheduled:
        reason = "queued"
    elif _foreground_chat_busy():
        reason = "chat_busy_deferred"
    else:
        reason = "already_inflight"
    return {
        "ok": True,
        "slug": slug,
        "scheduled": scheduled,
        "reason": reason,
        "embedding_reindex": get_embedding_reindex_status(),
    }


def _humanize_reindex_error(err: str) -> str:
    text = (err or "").strip()
    if not text:
        return text
    low = text.lower()
    if "no such file or directory" in low or "errno 2" in low:
        if "embedding model path not found" in low:
            return text
        if "mlx-embeddings" in low or "mlx stack" in low:
            return text
        return (
            f"{text} — check [utility] embedding_model_repo, MLX stack install, "
            "and that knowledge corpora still have synthesis.md or indexed sources"
        )
    if text in ("search_vectors_missing", "search_vectors_incomplete"):
        return "Search setup still in progress"
    if text == "chat_busy_deferred":
        return "Search setup will continue when chat is idle"
    return text


def _stamp_knowledge_thread_embed_identity(thread_id: str) -> None:
    """Persist MLX embed identity on a knowledge corpus after vectors are verified."""
    try:
        import session_search_db as ss

        stamp_conn = ss._safe_get_connection()
        if stamp_conn is None:
            return
        try:
            if ss.knowledge_semantic_index_complete(stamp_conn, thread_id):
                ss._stamp_thread_embed_identity_conn(stamp_conn, thread_id)
                stamp_conn.commit()
        finally:
            stamp_conn.close()
    except Exception as exc:
        _log.debug(
            "[embedding] stamp embed identity failed for %s: %s",
            thread_id,
            str(exc)[:120],
        )


def _finish_reindex_chat_busy_deferred(
    slug: str,
    *,
    emit_activity: bool,
    job_id: str,
    thread_id: str,
) -> dict[str, Any]:
    """Close an in-flight reindex activity when chat takes priority."""
    _emit_reindex_activity_completed(
        emit_activity=emit_activity,
        job_id=job_id,
        thread_id=thread_id,
        suffix="deferred for chat; retry scheduled",
    )
    return {"ok": False, "slug": slug, "error": "chat_busy_deferred"}


def _emit_reindex_activity_completed(
    *,
    emit_activity: bool,
    job_id: str,
    thread_id: str,
    suffix: str = "",
) -> None:
    if not emit_activity:
        return
    try:
        from utility_model import _activity_definition, _emit_activity

        op_name, op_def = _activity_definition("session_search_embedding")
        text = op_def["completed"]
        if suffix:
            text = f"{text} ({suffix})"
        _emit_activity(
            "completed",
            text,
            operation=op_name,
            job_id=job_id,
            thread_id=thread_id,
        )
    except Exception:
        pass


def _emit_reindex_activity_error(
    *,
    emit_activity: bool,
    job_id: str,
    thread_id: str,
    err: str,
) -> None:
    """Emit terminal error for an embedding reindex job (matches exception-path shape)."""
    if not emit_activity:
        return
    detail = _humanize_reindex_error(err)[:160]
    global _last_reindex_activity_err
    now = time.time()
    if (
        detail
        and detail == _last_reindex_activity_err[0]
        and (now - _last_reindex_activity_err[1]) < _ACTIVITY_ERR_DEDUPE_SEC
    ):
        return
    _last_reindex_activity_err = (detail, now)
    try:
        from utility_model import _activity_definition, _emit_activity

        op_name, op_def = _activity_definition("session_search_embedding")
        suffix = f": {detail}" if detail else ""
        _emit_activity(
            "error",
            f"{op_def['error']}{suffix}",
            operation=op_name,
            job_id=job_id,
            thread_id=thread_id,
        )
    except Exception:
        pass


def reindex_embeddings_for_corpus(slug: str, *, emit_activity: bool = False) -> dict[str, Any]:
    """Rebuild search vectors + knowledge graph embeddings for one corpus."""
    slug = (slug or "").strip()
    if not slug:
        return {"ok": False, "error": "missing_slug"}
    if embedding_work_should_defer():
        return {"ok": False, "slug": slug, "error": "chat_busy_deferred"}
    thread_id = f"knowledge:{slug}"
    job_id = f"embed-reindex:{slug}:{int(time.time() * 1000)}"
    if emit_activity:
        try:
            from utility_model import _activity_definition, _emit_activity

            op_name, op_def = _activity_definition("session_search_embedding")
            _emit_activity(
                "started",
                f"Upgrading search index ({slug})",
                operation=op_name,
                job_id=job_id,
                thread_id=thread_id,
            )
        except Exception:
            pass
    try:
        from tools.knowledge_tools import (
            _apply_chunk_meta_to_manifest_sources,
            _hydrate_raw_texts_for_indexing,
            _index_knowledge_in_search_db,
            _load_manifest,
            read_synthesis_text,
        )
        from utils.knowledge_graph import rebuild_knowledge_graph_for_corpus

        manifest = _load_manifest(slug) or {}
        label = str(manifest.get("label") or slug)
        sources = manifest.get("sources") if isinstance(manifest.get("sources"), list) else []
        synthesis_text = read_synthesis_text(slug)
        raw_texts = _hydrate_raw_texts_for_indexing(slug, sources, None)
        import session_search_db as ss

        had_complete_before = ss.knowledge_semantic_index_complete_for_thread(thread_id)
        prior_vec = ss.count_thread_vector_chunks(thread_id)
        search_current = _search_index_current_for_corpus(thread_id)
        if search_current:
            chunk_meta: dict[str, list[dict[str, Any]]] = {}
            _log.debug("[embedding] reindex %s: search index current; graph-only pass", slug)
        else:
            # _index_knowledge_in_search_db stages into a temp thread and swaps on success.
            chunk_meta = _index_knowledge_in_search_db(
                slug,
                label,
                synthesis_text,
                sources,
                raw_texts=raw_texts,
                ingest_profile=str(manifest.get("ingest_profile") or "docs"),
            )
        fts_rows = ss.count_thread_fts_messages(thread_id)
        if embedding_work_should_defer():
            return _finish_reindex_chat_busy_deferred(
                slug,
                emit_activity=emit_activity,
                job_id=job_id,
                thread_id=thread_id,
            )
        if (
            fts_rows > 0
            and not ss.knowledge_semantic_index_complete_for_thread(thread_id)
            and not search_current
        ):
            if ss.reindex_vectors_from_existing_fts(thread_id):
                _log.info(
                    "[embedding] reindex %s: recovered search vectors from existing FTS (%d rows)",
                    slug,
                    fts_rows,
                )
        had_material = bool(
            chunk_meta or (synthesis_text or "").strip() or sources
        )
        if fts_rows > 0 and not ss.knowledge_semantic_index_complete_for_thread(thread_id):
            err = (
                "search_vectors_missing"
                if ss.count_thread_vector_chunks(thread_id) <= 0
                else "search_vectors_incomplete"
            )
            now_vec = ss.count_thread_vector_chunks(thread_id)
            worsened = (had_complete_before and not ss.knowledge_semantic_index_complete_for_thread(thread_id)) or (
                now_vec < prior_vec
            )
            if emit_activity and worsened:
                _emit_reindex_activity_error(
                    emit_activity=True,
                    job_id=job_id,
                    thread_id=thread_id,
                    err=err,
                )
            elif emit_activity:
                if prior_vec <= 0 and now_vec <= 0:
                    _log.info(
                        "[embedding] reindex for %s still missing vectors (embedder unavailable or deferred)",
                        slug,
                    )
                _emit_reindex_activity_completed(
                    emit_activity=True,
                    job_id=job_id,
                    thread_id=thread_id,
                    suffix=f"{_humanize_reindex_error(err)}; retry scheduled",
                )
            return {"ok": False, "slug": slug, "error": err}
        if had_material and fts_rows <= 0:
            _log.warning(
                "[embedding] reindex for %s: manifest has material but search index is empty "
                "(chunk build returned nothing — try re-digest or check source files)",
                slug,
            )
            if emit_activity:
                _emit_reindex_activity_completed(
                    emit_activity=True,
                    job_id=job_id,
                    thread_id=thread_id,
                    suffix="no indexable chunks; re-digest corpus or check source files",
                )
            return {
                "ok": False,
                "slug": slug,
                "error": "no_indexable_chunks",
            }
        if embedding_work_should_defer():
            return _finish_reindex_chat_busy_deferred(
                slug,
                emit_activity=emit_activity,
                job_id=job_id,
                thread_id=thread_id,
            )
        search_index_ok = ss.knowledge_semantic_index_complete_for_thread(thread_id)
        _apply_chunk_meta_to_manifest_sources(slug, sources, chunk_meta)
        graph = rebuild_knowledge_graph_for_corpus(
            slug,
            label,
            synthesis_text,
            sources,
            raw_texts=raw_texts,
        )
        if graph.get("ok") is False:
            graph_err = str(graph.get("error") or "graph_embedding_failed")[:200]
            if search_index_ok:
                _stamp_knowledge_thread_embed_identity(thread_id)
            if emit_activity and not search_index_ok:
                _emit_reindex_activity_error(
                    emit_activity=True,
                    job_id=job_id,
                    thread_id=thread_id,
                    err=graph_err,
                )
            elif emit_activity and search_index_ok:
                _emit_reindex_activity_completed(
                    emit_activity=True,
                    job_id=job_id,
                    thread_id=thread_id,
                    suffix="graph rebuild deferred",
                )
            elif search_index_ok:
                _log.info(
                    "[embedding] search index ok for %s; graph rebuild failed: %s",
                    slug,
                    graph_err[:120],
                )
            return {
                "ok": False,
                "slug": slug,
                "error": graph_err,
                "search_index_ok": search_index_ok,
            }
        _stamp_knowledge_thread_embed_identity(thread_id)
        if emit_activity:
            _emit_reindex_activity_completed(
                emit_activity=True,
                job_id=job_id,
                thread_id=thread_id,
            )
        return {"ok": True, "slug": slug, "graph": graph}
    except Exception as exc:
        err = str(exc)[:200]
        _log.warning("[embedding] reindex failed for %s: %s", slug, err[:160])
        _emit_reindex_activity_error(
            emit_activity=emit_activity,
            job_id=job_id,
            thread_id=thread_id,
            err=err,
        )
        return {"ok": False, "slug": slug, "error": err}


def startup_embedding_identity_sweep(*, force: bool = False) -> None:
    """Background check: reindex corpora whose search index identity is stale."""
    global _startup_sweep_done
    if not force:
        with _LOAD_LOCK:
            if _startup_sweep_done:
                return
            _startup_sweep_done = True
    if not embedding_enabled():
        return

    def _runner() -> None:
        try:
            import session_search_db as ss

            cleared = ss.clear_orphan_reindex_staging_threads()
            if cleared:
                _log.info("[embedding] startup cleared %d orphan reindex staging thread(s)", cleared)
        except Exception as exc:
            _log.debug("[embedding] startup orphan staging cleanup failed: %s", str(exc)[:120])
        if embedding_enabled():
            try:
                prepare_embedding_model()
            except Exception:
                pass
        try:
            from tools.knowledge_tools import list_knowledge_corpora
        except Exception:
            return
        stale: list[str] = []
        for row in list_knowledge_corpora():
            slug = str(row.get("slug") or "").strip()
            if not slug:
                continue
            tid = f"knowledge:{slug}"
            if thread_needs_embedding_reindex(tid):
                stale.append(tid)
        if not stale:
            return
        if _foreground_chat_busy():
            _log.debug(
                "[embedding] startup sweep deferred — main chat active (%d corpora)",
                len(stale),
            )

            def _retry_sweep() -> None:
                startup_embedding_identity_sweep(force=True)

            threading.Timer(120.0, _retry_sweep).start()
            return
        if not load() and embedding_load_error():
            _log.info(
                "[embedding] startup sweep: MLX embedder unavailable (%s); "
                "using hash fallback for %d corpora",
                embedding_load_error()[:80],
                len(stale),
            )
        for tid in stale[:10]:
            maybe_schedule_embedding_reindex(tid)
            time.sleep(0.5)

    threading.Thread(target=_runner, daemon=True, name="embed-startup-sweep").start()


def prepare_embedding_model() -> bool:
    """Pre-load embedder (e.g. first run). Updates reindex status while downloading."""
    if not embedding_enabled():
        return False
    _set_reindex_status(state="preparing_model", corpora_pending=0, current_corpus="")
    ok = load()
    if ok:
        _set_reindex_status(state="idle", corpora_pending=0, current_corpus="")
    else:
        _set_reindex_status(state="idle", corpora_pending=0, current_corpus="")
    return ok
