"""Capability surface — which optional features are live right now.

Closes the honesty gap from session 102 item #2: previously, when
Ollama wasn't running or the embedding model wasn't pulled, semantic
search silently fell back to FTS5 keyword search without any visible
indicator. The UI (related-notes sidebar, Settings, search) claimed
"semantic search" while actually serving keyword matches.

This endpoint answers "is X actually available right now, and if not,
what does the user need to do to enable it?" — a single read-only
surface that the UI can render as a status badge.

Enumerated capabilities today:

- ``semantic_search`` — embeddings via Ollama (bge-m3 by default).
  ``ready=true`` means Ollama is reachable AND the model is pulled.
  ``reason`` carries the exact remediation string (``ollama serve``,
  ``ollama pull bge-m3``) when not ready.

- ``ocr`` — image text extraction via pytesseract. Optional extra
  (``pip install agntspace[ocr]``). ``ready=true`` when both the
  Python binding and the native Tesseract binary resolve.

- ``voice_transcription`` — Whisper-family ASR. Prefers the OpenAI
  Whisper API (``OPENAI_API_KEY`` in the credential store), falls back
  to faster-whisper if the ``[voice]`` extra is installed.

All probes are total (never raise) and non-blocking at the module
level — callers can poll this endpoint freely without worrying about
network timeouts killing the response (Ollama probe uses a short
connect timeout).
"""

from __future__ import annotations

import asyncio as _aio_caps
import importlib.util
import json
import logging
import re
import shutil
import time as _time_caps
from collections.abc import AsyncIterator
from typing import Any

import httpx
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

router = APIRouter(tags=["capabilities"])
log = logging.getLogger(__name__)


# Ollama's /api/pull is a long-running operation (multi-minute for
# big models on slow connections). We proxy the NDJSON stream back to
# the client verbatim so the UI can render a live progress bar.
# Defaults are shared with OllamaEmbedModel to avoid drift.
_DEFAULT_OLLAMA_BASE = "http://localhost:11434"

# Whitelist of model names we know how to pull. Restricts arbitrary
# strings from the request body — an attacker (or confused UI) can't
# make us pull `alpine:latest` or worse. The default (bge-m3) comes
# from vault/embeddings.py; adding a new model here requires a code
# change — deliberate, per Rule 12 "strategy in YAML" doesn't apply
# to safety whitelists.
_ALLOWED_MODELS = frozenset({  # arch:deterministic — security whitelist for embedding-model install endpoint, NOT user strategy
    "bge-m3",
    "nomic-embed-text",
    "mxbai-embed-large",
})

# Bounded pull duration — a stuck Ollama shouldn't keep the HTTP
# stream open indefinitely. 30 minutes covers bge-m3 (~1.2 GB) on a
# typical home connection; users on slow links can re-trigger.
_PULL_TIMEOUT_SECONDS = 30 * 60


class InstallModelRequest(BaseModel):
    """Install-model request body.

    ``model`` is validated against ``_ALLOWED_MODELS`` so the UI can't
    (accidentally or maliciously) make the server pull an arbitrary
    Ollama image. Defaults to bge-m3 which matches the default
    embedding model used by OllamaEmbedModel.
    """

    model: str = Field(default="bge-m3", min_length=1, max_length=64)


# ── Probes ───────────────────────────────────────────────────────────────


def _probe_semantic_search(request: Request) -> dict[str, Any]:
    """Probe Ollama embeddings health via ``VaultEmbeddings.health()``.

    Also reports how many chunks have been indexed so the UI can say
    "semantic search ready but 0 chunks indexed — run reindex" instead
    of the ambiguous "ready=true" with empty results.
    """
    emb = getattr(request.app.state, "vault_embeddings", None)
    if emb is None:
        return {
            "ready": False,
            "reason": "not_wired",
            "remediation": _remediation_for_reason("not_wired"),
            "model": "",
            "indexed_chunks": 0,
            "backend": "ollama",
        }

    try:
        health = emb.health()
    except Exception as exc:
        log.debug("Capabilities: embeddings.health() raised", exc_info=True)
        return {
            "ready": False,
            "reason": f"health_probe_failed:{type(exc).__name__}",
            "remediation": "Restart the server; if this persists, check the log for the real error.",
            "model": getattr(emb, "_model_name", ""),
            "indexed_chunks": 0,
            "backend": "ollama",
        }

    # The health dict from OllamaEmbedModel: {ready, reason, model,
    # base_url, identity}. We enrich it with the vector count and a
    # translated remediation string for the UI.
    remediation = _remediation_for_reason(
        health.get("reason", ""), model=health.get("model", ""),
    )

    return {
        "ready": bool(health.get("ready")),
        "reason": health.get("reason", ""),
        "remediation": remediation,
        "model": health.get("model", ""),
        "base_url": health.get("base_url", ""),
        "identity": health.get("identity", ""),
        "indexed_chunks": 0,  # Filled by async caller (can't block the probe)
        "backend": "ollama",
    }


def _remediation_for_reason(reason: str, *, model: str = "") -> str:
    """Translate a machine reason-string into a user-facing command hint.

    Kept local to this module so the UI doesn't have to parse reason
    strings. When we add a new failure mode to ``OllamaEmbedModel.health``,
    add the remediation here in the same PR.
    """
    if not reason or reason == "ok":
        return ""
    if reason.startswith("ollama_unreachable"):
        return "Start Ollama with `ollama serve`."
    if reason.startswith("model_not_pulled"):
        name = model or "bge-m3"
        return f"Install the model with `ollama pull {name}`."
    if reason.startswith("tags_http_"):
        return "Ollama responded with an error — check `ollama serve` logs."
    if reason == "tags_invalid_json":
        return "Ollama returned malformed JSON — check Ollama version / restart it."
    if reason == "not_wired":
        return "VaultEmbeddings isn't wired on this server instance."
    return ""


def _probe_ocr() -> dict[str, Any]:
    """Check if OCR (pytesseract + Tesseract binary) is usable."""
    # pytesseract Python binding
    if importlib.util.find_spec("pytesseract") is None:
        return {
            "ready": False,
            "reason": "pytesseract_missing",
            "remediation": "Install the optional extra: `pip install agntspace[ocr]`.",
            "backend": "pytesseract",
        }
    # Native Tesseract binary on PATH — pytesseract won't work without it.
    # We check `tesseract` by default; on Windows people sometimes install
    # to a hardcoded path that pytesseract.pytesseract.tesseract_cmd
    # points to, so we also try the package's resolved path before giving up.
    if shutil.which("tesseract"):
        return {
            "ready": True,
            "reason": "ok",
            "remediation": "",
            "backend": "pytesseract",
        }
    # Python binding present but binary not on PATH
    try:
        import pytesseract  # type: ignore[import-not-found]

        exe = getattr(pytesseract.pytesseract, "tesseract_cmd", "")
        from pathlib import Path

        if exe and Path(exe).is_file():
            return {
                "ready": True,
                "reason": "ok",
                "remediation": "",
                "backend": "pytesseract",
                "binary_path": str(exe),
            }
    except Exception:
        pass
    return {
        "ready": False,
        "reason": "tesseract_binary_missing",
        "remediation": "Install the Tesseract native binary "
                        "(https://github.com/UB-Mannheim/tesseract/wiki on Windows, "
                        "`brew install tesseract` on macOS, "
                        "`apt install tesseract-ocr` on Ubuntu).",
        "backend": "pytesseract",
    }


def _probe_voice() -> dict[str, Any]:
    """Voice transcription — prefers OpenAI Whisper API; falls back to faster-whisper."""
    # faster-whisper local backend (optional extra)
    faster_whisper_ok = importlib.util.find_spec("faster_whisper") is not None
    # openai-whisper (oldest local option, heavy dependency)
    openai_whisper_ok = importlib.util.find_spec("whisper") is not None
    # litellm is already a base dep — it can proxy to OpenAI's Whisper
    # API given an OPENAI_API_KEY. We can't probe for the key from here
    # without a request (credentials are scoped to the running app's
    # encrypted store), so we just report that the route exists.
    any_backend = faster_whisper_ok or openai_whisper_ok
    if any_backend:
        backend = "faster-whisper" if faster_whisper_ok else "openai-whisper"
        return {
            "ready": True,
            "reason": "ok",
            "remediation": "",
            "backend": backend,
        }
    return {
        "ready": False,
        "reason": "no_local_backend",
        "remediation": "Install the optional extra: `pip install agntspace[voice]` "
                        "(adds faster-whisper). OpenAI Whisper API also works when "
                        "`OPENAI_API_KEY` is configured under Settings → LLM.",
        "backend": "",
    }


# ── Route ────────────────────────────────────────────────────────────────


# /api/capabilities is polled by the Settings page on every render and
# every navigation. Each probe inside it makes a sync httpx call to
# Ollama (3s timeout) and importlib.find_spec scans for OCR / voice
# packages. Total cost: 1-3s per call, all on the event loop.
# Cache the bundle for 10s so a settings-page mount + scroll-spy +
# tab-switch share one probe round.
# Keyed by id(request.app) so different app instances (e.g. multiple
# pytest fixtures, multi-tenant servers) get independent caches.
_CAPABILITIES_CACHE: dict[int, dict[str, Any]] = {}
# 60-second TTL: capability state (Ollama reachable, OCR/voice
# packages installed) is essentially immutable on the timescale of a
# session. The Settings page reads this on every render and tab
# switch. Prior 10s TTL meant most reads paid the full Ollama probe
# cost; 60s eliminates the thundering-herd cold-cache hit on dashboard
# mount. The install endpoint resets the cache explicitly.
_CAPABILITIES_TTL = 60.0


def _build_capabilities_blocking(request: Request) -> dict[str, Any]:
    """Sync helper — runs every probe in one thread off the event loop."""
    return {
        "semantic_search": _probe_semantic_search(request),
        "ocr": _probe_ocr(),
        "voice_transcription": _probe_voice(),
    }


def _reset_capabilities_cache() -> None:
    """Clear the per-app capability cache. Used by tests + after install/restart actions."""
    _CAPABILITIES_CACHE.clear()


@router.get("/capabilities")
async def get_capabilities(request: Request) -> dict[str, Any]:
    """Return the live status of every optional capability.

    Shape::

        {
          "semantic_search": {
            "ready": bool,
            "reason": str,            # machine tag (ok / ollama_unreachable / ...)
            "remediation": str,       # translated user-facing hint
            "model": "bge-m3",
            "base_url": "http://localhost:11434",
            "identity": "ollama:bge-m3",
            "indexed_chunks": int,
            "backend": "ollama"
          },
          "ocr": { ready, reason, remediation, backend, binary_path? },
          "voice_transcription": { ready, reason, remediation, backend }
        }

    This is the SINGLE source of truth the UI reads to render capability
    badges. Pure read — never mutates state, safe to poll.

    Blocking probes (sync httpx to Ollama, importlib package scans) run
    off the asyncio loop via to_thread. Cached 10s so dashboard polls
    don't pay the probe cost on every request.
    """
    now = _time_caps.monotonic()
    cache_key = id(request.app)
    cached = _CAPABILITIES_CACHE.get(cache_key)
    if cached and (now - cached["ts"]) < _CAPABILITIES_TTL and cached.get("data") is not None:
        out = dict(cached["data"])
    else:
        try:
            out = await _aio_caps.to_thread(_build_capabilities_blocking, request)
        except Exception:
            out = {
                "semantic_search": {"ready": False, "reason": "probe_failed", "remediation": "", "model": "", "indexed_chunks": 0, "backend": "ollama"},
                "ocr": {"ready": False, "reason": "probe_failed", "remediation": "", "backend": "tesseract"},
                "voice_transcription": {"ready": False, "reason": "probe_failed", "remediation": "", "backend": "none"},
            }
        _CAPABILITIES_CACHE[cache_key] = {"ts": now, "data": out}

    # Fill in indexed_chunks fresh on every call — the count changes as
    # embeddings index grows and the DB read is properly async.
    semantic = dict(out.get("semantic_search", {}))
    db = getattr(request.app.state, "db", None)
    if db is not None:
        try:
            cursor = await db.execute(
                "SELECT COUNT(*) FROM vault_embeddings"
            )
            row = await cursor.fetchone()
            semantic["indexed_chunks"] = int(row[0]) if row else 0
        except Exception:
            # Table may not exist yet (fresh vault, no init_schema call)
            semantic["indexed_chunks"] = 0

    return {
        "semantic_search": semantic,
        "ocr": out.get("ocr", {}),
        "voice_transcription": out.get("voice_transcription", {}),
    }


# ── Install surface (session 102 item #3) ────────────────────────────────


def _resolve_ollama_base(request: Request) -> str:
    """Pick the Ollama base URL the install endpoint should hit.

    Prefers the URL the live VaultEmbeddings adapter is configured with
    (so install matches the probe) and falls back to the module default.
    Keeps semantic_search.base_url and the install target in lockstep.
    """
    emb = getattr(request.app.state, "vault_embeddings", None)
    if emb is not None:
        url = getattr(emb, "_base_url", "")
        if isinstance(url, str) and url:
            return url.rstrip("/")
    return _DEFAULT_OLLAMA_BASE


async def _stream_ollama_pull(
    base_url: str, model: str,
) -> AsyncIterator[bytes]:
    """Stream Ollama's /api/pull response as NDJSON bytes.

    Ollama emits one JSON object per status transition:
      {"status": "pulling manifest"}
      {"status": "downloading", "digest": "...", "total": N, "completed": M}
      ...
      {"status": "verifying sha256 digest"}
      {"status": "success"}

    We pass every line through to the client verbatim so the UI can
    parse the same progress information Ollama's CLI shows. On the
    final event we emit a synthetic ``{"status":"ok","done":true}``
    so the client knows the stream ended cleanly rather than being
    cut off mid-pull (which it would otherwise look like).
    """
    async with httpx.AsyncClient(timeout=_PULL_TIMEOUT_SECONDS) as client:
        try:
            async with client.stream(
                "POST",
                f"{base_url}/api/pull",
                json={"model": model, "stream": True},
            ) as resp:
                if resp.status_code != 200:
                    text = (await resp.aread()).decode("utf-8", errors="replace")[:200]
                    yield _err_line(
                        f"Ollama /api/pull returned {resp.status_code}: {text}",
                    )
                    return
                async for chunk in resp.aiter_raw():
                    yield chunk
        except httpx.ConnectError as exc:
            # Most common non-happy path: user stopped Ollama between
            # probe and install click, or never started it. Give them
            # the exact command.
            yield _err_line(
                f"Ollama is not reachable at {base_url}. Start it with `ollama serve`. ({exc})",
            )
            return
        except httpx.HTTPError as exc:
            yield _err_line(f"Ollama request failed: {exc}")
            return
        except Exception as exc:  # noqa: BLE001 — we never want to leak a stack trace to the stream
            log.exception("Unexpected error while pulling model %s", model)
            yield _err_line(f"Unexpected error: {type(exc).__name__}: {exc}")
            return

    # Synthetic terminator so the UI knows the stream ended cleanly.
    yield _ok_line("completed")


def _err_line(message: str) -> bytes:
    """Format an NDJSON error line the UI can render inline."""
    return (json.dumps({"status": "error", "error": message}) + "\n").encode("utf-8")


def _ok_line(reason: str) -> bytes:
    """Format an NDJSON terminator the UI can detect."""
    return (json.dumps({"status": "done", "reason": reason}) + "\n").encode("utf-8")


def _validate_model_name(model: str) -> str:
    """Reject arbitrary model strings — only the whitelist installs.

    Raises HTTPException(400) on violation so FastAPI surfaces a clean
    error to the caller. Model name itself is also length/character
    sanity-checked so a `../../etc/passwd` style string can't reach
    the Ollama URL (paranoia — Ollama's `/api/pull` treats this as
    JSON body, not URL path, but defense in depth).
    """
    if not isinstance(model, str) or not model:
        raise HTTPException(status_code=400, detail="model must be a non-empty string")
    # Only allow simple model names: lowercase, digits, hyphens, periods,
    # and a single colon for the optional tag (e.g. "bge-m3:latest").
    if not re.match(r"^[a-z0-9][a-z0-9.\-]*(:[a-z0-9.\-]+)?$", model):
        raise HTTPException(
            status_code=400,
            detail="model name must be lowercase alphanumeric with hyphens/periods, optionally suffixed `:tag`",
        )
    # Strip any :tag for the whitelist check so "bge-m3:latest" still passes.
    base = model.split(":", 1)[0]
    if base not in _ALLOWED_MODELS:
        raise HTTPException(
            status_code=400,
            detail=f"model `{base}` not in install whitelist; allowed: {sorted(_ALLOWED_MODELS)}",
        )
    return model


@router.post("/capabilities/semantic_search/install")
async def install_semantic_model(
    request: Request, body: InstallModelRequest,
) -> StreamingResponse:
    """Pull the embedding model via Ollama, streaming NDJSON progress.

    Closes session 102 item #3: before this route, a user without
    ``bge-m3`` pulled had to open a terminal and type
    ``ollama pull bge-m3`` after reading the remediation hint. This
    endpoint lets them click a button in Settings → Capabilities and
    watch progress in the UI.

    **Shape of the streamed NDJSON** (one JSON object per line):

    - Normal status from Ollama: ``{"status": "downloading", "digest":
      "sha256:...", "total": 123456, "completed": 50000}`` — the UI
      renders progress as ``completed / total`` bytes.
    - Terminal success: ``{"status": "done", "reason": "completed"}``
      (synthetic — generated after Ollama's ``{"status": "success"}``
      so the UI can distinguish "stream ended cleanly" from "stream
      cut off mid-pull").
    - Any failure: ``{"status": "error", "error": "<message>"}`` —
      the UI surfaces this to the user with the remediation command.

    **Model whitelist**: the request body's ``model`` field is
    validated against ``_ALLOWED_MODELS``. Arbitrary strings are
    rejected with HTTP 400 so the UI can't make us pull e.g.
    ``some-attacker-registry/malicious:latest``.

    **Idempotent**: if the model is already installed, Ollama returns
    immediately with ``{"status": "success"}`` — safe to call multiple
    times without redownloading.
    """
    model = _validate_model_name(body.model)
    base_url = _resolve_ollama_base(request)

    # ActivityLogger event so the install shows up in the decisions
    # timeline — Rule 13: user-visible actions must be observable.
    activity = getattr(request.app.state, "activity_logger", None)
    if activity is not None:
        try:
            activity.log(
                category="system",
                action="embedding_model_install_started",
                summary=f"Pulling embedding model `{model}` from Ollama",
                details={"model": model, "base_url": base_url},
                importance="medium",
            )
        except Exception:
            log.debug("ActivityLogger.log failed on install_started", exc_info=True)

    return StreamingResponse(
        _stream_ollama_pull(base_url, model),
        media_type="application/x-ndjson",
        headers={
            # Disable proxy buffering so progress arrives in real time.
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
        },
    )
