"""Vector embeddings for semantic vault search — Ollama-native.

Embeddings run locally via Ollama's ``POST /api/embed`` endpoint. No fastembed,
no Hugging Face Hub, no tokenizers — only ``httpx`` (BSD-3) and the Python
standard library. Vectors are stored in SQLite alongside the FTS5 index.

The default model is ``bge-m3`` (multilingual, 1024-dim). The user pulls it
via ``ollama pull bge-m3`` (one-time, ~1.2 GB on disk) — the server never
downloads model weights autonomously.

Hybrid search: 70% vector similarity + 30% FTS5 keyword. Both flavors
fall back to FTS5-only when Ollama isn't reachable, surfaced as
``EmbeddingsUnavailable`` so callers can render a clear capability state.
"""

from __future__ import annotations

import hashlib
import logging
import math
import struct
from collections.abc import Iterable
from datetime import UTC, datetime

import aiosqlite
import httpx

from agntspace.vault.reader import VaultReader

log = logging.getLogger(__name__)

# ── Defaults ──────────────────────────────────────────────────────────────
# Ollama's default listen address. Matches llm/ollama_discovery.OLLAMA_BASE
# but we avoid the cross-module import to keep embeddings independent of
# the LLM provider layer.
DEFAULT_OLLAMA_BASE = "http://localhost:11434"

# bge-m3 is the multilingual SOTA embedding model available from Ollama
# via `ollama pull bge-m3`. 1024-dim, supports 100+ languages, handles
# cross-language queries natively. Users can override via
# AGNT_EMBED_MODEL env var or by passing model= to VaultEmbeddings.
DEFAULT_EMBED_MODEL = "bge-m3"

# Stable identifier stored in vault_embedding_meta so a model switch
# (e.g. bge-m3 → nomic-embed-text) wipes old vectors on next boot.
def _model_identity(model_name: str) -> str:
    return f"ollama:{model_name}"


CHUNK_SIZE = 500  # arch:deterministic -- bge-m3 has a 512-token sequence ceiling and tokenization expands ~1.3x; 500 chars is the structurally safe upper bound for text that fits in one embedding pass without truncation, not a user-tunable preference
CHUNK_OVERLAP = 50  # arch:deterministic -- 10% overlap is the empirically standard value for sliding-window text embedding to preserve cross-chunk context, fixed by the embedder's input window not by user choice

# ── Activity counters (read by the /api/llm/local-stack UI) ──
# Lightweight per-process telemetry so the Local Stack panel can show the
# "⚡ Active" pulse on bge-m3 when embeddings are flowing. Not persisted —
# these reset on every server restart, which is fine because the UI refreshes
# live while the server is up. Writes are single-assignment + int increment,
# both atomic under the GIL, so no lock is needed.
_total_embed_calls: int = 0
_last_embed_call_iso: str = ""
_last_embed_model: str = ""
_last_embed_batch_size: int = 0


class EmbeddingsUnavailable(RuntimeError):
    """Raised when Ollama isn't reachable OR the embedding model isn't pulled.

    Distinct from generic network errors so API routes can degrade cleanly
    to FTS5 instead of emitting 5xx responses. The message carries the
    exact remediation (``ollama pull bge-m3`` / ``ollama serve``) so the
    capability-indicator UI surface can render an actionable hint.
    """


# ── HTTP adapter ──────────────────────────────────────────────────────────


class OllamaEmbedModel:
    """Thin HTTP adapter matching the ``model.embed(texts)`` contract that
    ``VaultEmbeddings``, ``KnowledgeGraph``, ``EntityRegistry``, and the
    skill loader all rely on.

    The adapter is synchronous because every caller already either runs
    in a sync executor OR uses list(model.embed(...)) eagerly. Switching
    to async would ripple into three layers of calling code for no real
    benefit (embeddings are CPU-bound on the Ollama side; we're just
    doing a blocking HTTP wait).
    """

    def __init__(
        self,
        model: str = DEFAULT_EMBED_MODEL,
        base_url: str = DEFAULT_OLLAMA_BASE,
        *,
        # 60s was a cold-start allowance when the embedding pipeline was a
        # background warm-up task. Now it runs inline on entity resolution
        # (hot path, called from ingest/chat/graph-populate) — a 40s block
        # on bge-m3 cold-start freezes the event loop and makes health
        # probes time out. 8s is still generous for a warm call (~50ms) and
        # short enough that one cold-start trips the EntityRegistry circuit
        # breaker before the server looks hung. Re-tune in config if a
        # heavier model legitimately needs longer warm-up.
        timeout: float = 8.0,
    ) -> None:
        self.model = model
        self.base_url = base_url.rstrip("/")
        self._timeout = timeout
        # Lazy-built client so tests that never touch the network never
        # open a real socket. Recreated on each call is cheap enough but
        # we cache per-instance for connection reuse on the batch path.
        self._client: httpx.Client | None = None

    def _http(self) -> httpx.Client:
        if self._client is None:
            self._client = httpx.Client(timeout=self._timeout)
        return self._client

    def close(self) -> None:
        if self._client is not None:
            try:
                self._client.close()
            except Exception:
                pass
            self._client = None

    def __del__(self) -> None:  # noqa: D105
        self.close()

    def embed(self, texts: Iterable[str]) -> list[list[float]]:
        """Embed a batch of texts. Returns one ``list[float]`` per input.

        Raises:
            EmbeddingsUnavailable: if Ollama isn't reachable or the model
                isn't pulled. Caller should degrade to FTS5-only.
        """
        # Normalize to list so we can len() + iterate deterministically.
        batch = [t if t else " " for t in texts]
        if not batch:
            return []
        # Record per-process activity so the Local Stack UI can show
        # "⚡ Active" on the model that's doing work. Bumped before the
        # network call so a long-running or failing call is visible to
        # the UI (helps diagnose slow Ollama — the exact problem we hit
        # earlier today with bge-m3 cold-start).
        global _total_embed_calls, _last_embed_call_iso, _last_embed_model, _last_embed_batch_size
        _total_embed_calls += 1
        _last_embed_model = self.model
        _last_embed_batch_size = len(batch)
        try:
            from datetime import UTC, datetime
            _last_embed_call_iso = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
        except Exception:
            pass
        try:
            resp = self._http().post(
                f"{self.base_url}/api/embed",
                json={"model": self.model, "input": batch},
            )
        except httpx.ConnectError as exc:
            raise EmbeddingsUnavailable(
                f"Ollama is not reachable at {self.base_url}. "
                f"Start it with `ollama serve`."
            ) from exc
        except httpx.HTTPError as exc:
            raise EmbeddingsUnavailable(
                f"Ollama request failed: {exc}"
            ) from exc

        if resp.status_code == 404:
            # Ollama returns 404 when the model isn't pulled locally.
            raise EmbeddingsUnavailable(
                f"Embedding model `{self.model}` not installed in Ollama. "
                f"Pull it with `ollama pull {self.model}`."
            )
        if resp.status_code >= 400:
            raise EmbeddingsUnavailable(
                f"Ollama /api/embed returned {resp.status_code}: "
                f"{resp.text[:200]}"
            )

        try:
            data = resp.json()
        except ValueError as exc:
            raise EmbeddingsUnavailable(
                f"Ollama /api/embed returned non-JSON: {resp.text[:200]}"
            ) from exc

        vectors = data.get("embeddings")
        if vectors is None:
            # Single-input API may return "embedding" (singular)
            single = data.get("embedding")
            if single is not None:
                vectors = [single]
        if not isinstance(vectors, list) or len(vectors) != len(batch):
            raise EmbeddingsUnavailable(
                f"Ollama /api/embed returned unexpected shape: "
                f"expected {len(batch)} vectors, got {type(vectors).__name__}"
            )
        # Force float — Ollama returns JSON numbers which may be int 0
        # on degenerate inputs.
        return [[float(x) for x in vec] for vec in vectors]

    def health(self) -> dict:
        """Quick probe: is Ollama up AND is the model installed.

        Returns ``{"ready": bool, "reason": str, "model": str, "base_url": str}``.
        Never raises. Used by the capability-indicator surface so the UI
        can render "semantic search enabled / falling back to keyword".
        """
        base = {"ready": False, "model": self.model, "base_url": self.base_url}
        try:
            resp = self._http().get(f"{self.base_url}/api/tags", timeout=2.0)
        except httpx.HTTPError as exc:
            return {**base, "reason": f"ollama_unreachable:{type(exc).__name__}"}
        if resp.status_code != 200:
            return {**base, "reason": f"tags_http_{resp.status_code}"}
        try:
            installed = {
                (m.get("name") or m.get("model") or "").lower()
                for m in resp.json().get("models", [])
            }
        except ValueError:
            return {**base, "reason": "tags_invalid_json"}
        # Match by bare model name OR the :latest tag Ollama often appends.
        wanted = self.model.lower()
        if wanted in installed or f"{wanted}:latest" in installed or any(
            n.startswith(f"{wanted}:") for n in installed
        ):
            return {**base, "ready": True, "reason": "ok"}
        return {**base, "reason": f"model_not_pulled:{self.model}"}


# ── Pure-Python vector math (no numpy) ────────────────────────────────────


def _cosine(u: list[float], v: list[float]) -> float:
    """Cosine similarity for two equal-length vectors.

    Returns 0.0 when either norm is zero. Pure Python so we don't pull
    numpy (BSD but avoided to keep the embeddings layer MIT-only).
    For 1024-dim vectors over a few thousand entities this runs in well
    under a second on CPython; switch to numpy only if profiling says so.
    """
    if not u or not v or len(u) != len(v):
        return 0.0
    dot = 0.0
    nu = 0.0
    nv = 0.0
    for a, b in zip(u, v):
        dot += a * b
        nu += a * a
        nv += b * b
    if nu <= 0.0 or nv <= 0.0:
        return 0.0
    return dot / (math.sqrt(nu) * math.sqrt(nv))


def _vec_to_bytes(vec: list[float]) -> bytes:
    """Pack vector as little-endian float32. Matches numpy's .tobytes()
    layout so previously-stored vectors (if any) decode identically."""
    return struct.pack(f"<{len(vec)}f", *vec)


def _bytes_to_vec(data: bytes) -> list[float]:
    """Unpack float32 LE bytes into a list[float]. Silently returns [] on
    a corrupt row rather than crashing a search."""
    n, rem = divmod(len(data), 4)
    if rem != 0 or n == 0:
        return []
    return list(struct.unpack(f"<{n}f", data))


# ── Public service ────────────────────────────────────────────────────────


class VaultEmbeddings:
    """Semantic search over vault using Ollama-hosted embeddings."""

    def __init__(
        self,
        db: aiosqlite.Connection,
        vault: VaultReader,
        exclude_dirs: list[str] | None = None,
        *,
        model: str = DEFAULT_EMBED_MODEL,
        base_url: str = DEFAULT_OLLAMA_BASE,
    ) -> None:
        self._db = db
        self._vault = vault
        self._exclude_dirs = set(exclude_dirs or [])
        self._model_name = model
        self._base_url = base_url
        self._model: OllamaEmbedModel | None = None  # lazy

    def _get_model(self) -> OllamaEmbedModel:
        """Return (and cache) the Ollama adapter. Does NOT probe the
        network — the first ``.embed(...)`` call will surface
        ``EmbeddingsUnavailable`` if Ollama is down or the model isn't
        pulled. Keeping this method total and cheap preserves the
        shape SkillLoader / KnowledgeGraph / EntityRegistry rely on
        (``embedder._get_model()`` returning an object with ``.embed()``).
        """
        if self._model is None:
            self._model = OllamaEmbedModel(
                model=self._model_name, base_url=self._base_url,
            )
        return self._model

    def health(self) -> dict:
        """Expose the Ollama probe for the capability-indicator surface.

        Returns the same shape as ``OllamaEmbedModel.health()`` plus an
        ``identity`` field the UI can use to distinguish runtime backends
        (stable across model swaps within the same provider family).
        """
        h = self._get_model().health()
        return {**h, "identity": _model_identity(self._model_name)}

    async def init_schema(self) -> None:
        """Create embedding tables if absent, and wipe vectors when the
        backing model changes so stale (wrong-dim or wrong-space) vectors
        don't pollute search results."""
        await self._db.executescript("""
            CREATE TABLE IF NOT EXISTS vault_embeddings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                path TEXT NOT NULL,
                chunk_index INTEGER NOT NULL,
                chunk_text TEXT NOT NULL,
                embedding BLOB NOT NULL,
                content_hash TEXT NOT NULL,
                indexed_at TEXT NOT NULL,
                UNIQUE(path, chunk_index)
            );
            CREATE INDEX IF NOT EXISTS idx_emb_path ON vault_embeddings(path);
            CREATE TABLE IF NOT EXISTS vault_embedding_meta (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            );
        """)
        await self._db.commit()

        cursor = await self._db.execute(
            "SELECT value FROM vault_embedding_meta WHERE key = 'model'"
        )
        row = await cursor.fetchone()
        stored = row[0] if row else None
        current = _model_identity(self._model_name)
        if stored and stored != current:
            log.warning(
                "Embedding model changed %s → %s — clearing stale vectors (will re-index)",
                stored, current,
            )
            await self._db.execute("DELETE FROM vault_embeddings")
        await self._db.execute(
            "INSERT OR REPLACE INTO vault_embedding_meta (key, value) VALUES ('model', ?)",
            (current,),
        )
        await self._db.commit()

    async def index_vault(self) -> int:
        """Index all vault markdown files with embeddings. Returns count.

        Raises ``EmbeddingsUnavailable`` when Ollama is off so callers
        (e.g. /api/search/reindex) can return a clear error instead of
        silently writing no rows.
        """
        model = self._get_model()
        files = self._vault.list_files(pattern="**/*.md")
        indexed = 0

        for file_path in files:
            if self._vault.paths:
                relative = self._vault.paths.to_logical(file_path)
            else:
                relative = str(file_path.relative_to(self._vault.vault_dir)).replace("\\", "/")

            if any(relative.startswith(f"{d}/") for d in self._exclude_dirs):
                continue

            try:
                doc = self._vault.read(relative)
            except Exception:
                continue

            content_hash = hashlib.md5(doc.raw.encode()).hexdigest()

            cursor = await self._db.execute(
                "SELECT content_hash FROM vault_embeddings WHERE path = ? LIMIT 1",
                (relative,),
            )
            row = await cursor.fetchone()
            if row and row[0] == content_hash:
                continue

            await self._db.execute("DELETE FROM vault_embeddings WHERE path = ?", (relative,))

            chunks = self._chunk_text(doc.content)
            if not chunks:
                continue

            embeddings = model.embed(chunks)

            now = datetime.now(UTC).isoformat()
            for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
                emb_bytes = _vec_to_bytes(embedding)
                await self._db.execute(
                    "INSERT INTO vault_embeddings (path, chunk_index, chunk_text, embedding, content_hash, indexed_at) "
                    "VALUES (?, ?, ?, ?, ?, ?)",
                    (relative, i, chunk, emb_bytes, content_hash, now),
                )

            indexed += 1

        await self._db.commit()
        log.info("Vault embeddings indexed: %d files", indexed)
        return indexed

    async def search(self, query: str, limit: int = 10) -> list[dict]:
        """Semantic search via cosine similarity over stored vectors.

        Raises ``EmbeddingsUnavailable`` when Ollama is off (the caller
        in /api/search/semantic catches this and degrades to FTS5).
        """
        if not query.strip():
            return []

        model = self._get_model()
        query_vectors = model.embed([query])
        if not query_vectors:
            return []
        q = query_vectors[0]

        cursor = await self._db.execute(
            "SELECT path, chunk_index, chunk_text, embedding FROM vault_embeddings"
        )
        rows = await cursor.fetchall()
        if not rows:
            return []

        results: list[dict] = []
        for row in rows:
            path, chunk_idx, chunk_text, emb_bytes = row
            vec = _bytes_to_vec(emb_bytes)
            if not vec:
                continue
            score = _cosine(q, vec)
            if score == 0.0:
                continue
            results.append({
                "path": path,
                "chunk_index": chunk_idx,
                "text": chunk_text,
                "score": float(score),
            })

        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:limit]

    async def hybrid_search(
        self, query: str, limit: int = 10, vector_weight: float = 0.7,
    ) -> list[dict]:
        """Hybrid: 70% semantic + 30% keyword FTS5."""
        from agntspace.vault.search import VaultSearch

        vector_results = await self.search(query, limit=limit * 2)

        fts_search = VaultSearch(self._db, self._vault)
        keyword_results = await fts_search.search(query, limit=limit * 2)

        max_vec = max((r["score"] for r in vector_results), default=1.0)
        max_kw = max((abs(r.get("score", 0)) for r in keyword_results), default=1.0)

        merged: dict[str, dict] = {}
        for r in vector_results:
            key = f"{r['path']}:{r.get('chunk_index', 0)}"
            merged[key] = {
                "path": r["path"],
                "text": r["text"],
                "vector_score": r["score"] / max_vec if max_vec else 0,
                "keyword_score": 0,
            }

        for r in keyword_results:
            key = f"{r['path']}:0"
            if key in merged:
                merged[key]["keyword_score"] = (
                    abs(r.get("score", 0)) / max_kw if max_kw else 0
                )
            else:
                merged[key] = {
                    "path": r["path"],
                    "text": r.get("snippet", ""),
                    "vector_score": 0,
                    "keyword_score": abs(r.get("score", 0)) / max_kw if max_kw else 0,
                }

        results = []
        for item in merged.values():
            hybrid = (
                vector_weight * item["vector_score"]
                + (1 - vector_weight) * item["keyword_score"]
            )
            results.append({
                "path": item["path"],
                "text": item["text"],
                "score": hybrid,
                "vector_score": item["vector_score"],
                "keyword_score": item["keyword_score"],
            })

        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:limit]

    @staticmethod
    def _chunk_text(text: str) -> list[str]:
        """Split text into overlapping chunks."""
        if len(text) <= CHUNK_SIZE:
            return [text] if text.strip() else []
        chunks = []
        start = 0
        while start < len(text):
            end = start + CHUNK_SIZE
            chunk = text[start:end]
            if chunk.strip():
                chunks.append(chunk.strip())
            start = end - CHUNK_OVERLAP
        return chunks
