"""ObsiGate — Semantic search: embeddings, vector store and hybrid retrieval.

This module adds a *semantic* layer on top of the existing TF-IDF search. Each
document is split into overlapping chunks, each chunk is converted into a dense
vector, and queries are matched by cosine similarity. Results are combined with
the lexical ranking through Reciprocal Rank Fusion (RRF).

Design goals
------------
* **Zero mandatory dependency.** ``sentence-transformers`` (local model),
  ``numpy`` and ``faiss`` are *optional*. They are imported lazily and, when
  missing, the module falls back to a deterministic pure-Python hashing embedder
  and a pure-Python cosine store. The feature therefore degrades gracefully and
  the default CI (which only installs ``backend/requirements.txt``) keeps working.
* **Plug-in providers.** Embeddings can come from the local
  ``all-MiniLM-L6-v2`` model, from an OpenAI-compatible ``/embeddings`` endpoint
  (configured via env vars), or from the deterministic fallback.
* **Incremental.** The index is updated document-by-document from the indexer
  change hook (file watcher + API mutations), never rebuilt on each search.

Optional extras are listed in ``backend/requirements-semantic.txt``.
"""

from __future__ import annotations

import hashlib
import logging
import math
import os
import re
import threading
from abc import ABC, abstractmethod
from collections import Counter
from itertools import pairwise
from typing import Any

logger = logging.getLogger("obsigate.semantic")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
EMBEDDING_DIM = 384               # all-MiniLM-L6-v2 output dimension
CHUNK_TOKENS = 512                # target chunk size (whitespace tokens)
CHUNK_OVERLAP_TOKENS = 64         # overlap between consecutive chunks
DEFAULT_TOP_K = 200               # max documents returned by a semantic query
RRF_K = 60                        # Reciprocal Rank Fusion smoothing constant
MAX_QUERY_CHARS = 2000            # guard against pathological queries

_WORD_RE = re.compile(r"[\w]+", re.UNICODE)


# ---------------------------------------------------------------------------
# Tokenization / chunking
# ---------------------------------------------------------------------------
def _simple_tokens(text: str) -> list[str]:
    """Split *text* into lowercase word tokens (keeps accents)."""
    return _WORD_RE.findall(text.lower())


def chunk_text(
    text: str,
    chunk_tokens: int = CHUNK_TOKENS,
    overlap: int = CHUNK_OVERLAP_TOKENS,
) -> list[str]:
    """Split *text* into overlapping windows of roughly *chunk_tokens* words.

    Args:
        text: Raw document text.
        chunk_tokens: Target number of whitespace tokens per chunk.
        overlap: Number of tokens shared by two consecutive chunks.

    Returns:
        A list of chunk strings. Empty input yields an empty list.
    """
    if not text or not text.strip():
        return []
    if chunk_tokens <= 0:
        chunk_tokens = CHUNK_TOKENS
    overlap = max(0, min(overlap, chunk_tokens - 1))

    words = text.split()
    if len(words) <= chunk_tokens:
        return [" ".join(words)]

    step = max(1, chunk_tokens - overlap)
    chunks: list[str] = []
    for start in range(0, len(words), step):
        window = words[start:start + chunk_tokens]
        if not window:
            break
        chunks.append(" ".join(window))
        if start + chunk_tokens >= len(words):
            break
    return chunks


# ---------------------------------------------------------------------------
# Embedding providers
# ---------------------------------------------------------------------------
class EmbeddingProvider(ABC):
    """Base class for embedding backends."""

    name: str = "base"

    def __init__(self, dimension: int = EMBEDDING_DIM) -> None:
        self.dimension = dimension

    @abstractmethod
    def encode(self, texts: list[str]) -> list[list[float]]:
        """Return one L2-normalized vector per input text."""

    def encode_one(self, text: str) -> list[float]:
        """Convenience wrapper returning the vector for a single text."""
        vectors = self.encode([text])
        return vectors[0] if vectors else [0.0] * self.dimension


class HashEmbeddingProvider(EmbeddingProvider):
    """Deterministic, dependency-free hashing embedder.

    This is a *lexical* fallback: it hashes word unigrams, word bigrams and
    character trigrams into fixed-size signed buckets (the "hashing trick"),
    then L2-normalizes the result. It captures shared vocabulary and
    morphological variants (``backup``/``backups``), so it already improves
    recall over exact TF-IDF matching, but it does not understand synonyms the
    way a real transformer model does.
    """

    name = "hash"

    def _add_feature(self, vec: list[float], key: str, weight: float) -> None:
        digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
        h = int.from_bytes(digest, "big")
        idx = h % self.dimension
        sign = 1.0 if (h >> 63) & 1 else -1.0
        vec[idx] += sign * weight

    def _encode_one(self, text: str) -> list[float]:
        vec = [0.0] * self.dimension
        tokens = _simple_tokens(text)
        if not tokens:
            return vec

        tf = Counter(tokens)
        for token, count in tf.items():
            weight = 1.0 + math.log(count)
            self._add_feature(vec, "w:" + token, weight)
            for gram in _char_ngrams(token, 3):
                self._add_feature(vec, "g:" + gram, weight * 0.5)

        for first, second in pairwise(tokens):
            self._add_feature(vec, "b:" + first + "_" + second, 0.5)

        norm = math.sqrt(sum(v * v for v in vec))
        if norm > 0.0:
            vec = [v / norm for v in vec]
        return vec

    def encode(self, texts: list[str]) -> list[list[float]]:
        return [self._encode_one(t or "") for t in texts]


def _char_ngrams(token: str, n: int) -> list[str]:
    """Return padded character n-grams for *token* (bounded to avoid blow-up)."""
    if len(token) < n:
        return [token]
    if len(token) > 24:
        token = token[:24]
    return [token[i:i + n] for i in range(len(token) - n + 1)]


class SentenceTransformerProvider(EmbeddingProvider):
    """Local ``all-MiniLM-L6-v2`` embeddings via ``sentence-transformers``."""

    name = "sentence-transformers"

    def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
        super().__init__(EMBEDDING_DIM)
        self.model_name = model_name
        self._model: Any | None = None

    @staticmethod
    def is_available() -> bool:
        try:
            import sentence_transformers  # noqa: F401
        except Exception:
            return False
        return True

    def _get_model(self) -> Any:
        if self._model is None:
            from sentence_transformers import SentenceTransformer

            self._model = SentenceTransformer(self.model_name)
        return self._model

    def encode(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        model = self._get_model()
        vectors = model.encode(texts, normalize_embeddings=True)
        return [[float(x) for x in vec] for vec in vectors]


class RemoteEmbeddingProvider(EmbeddingProvider):
    """OpenAI-compatible ``/embeddings`` endpoint (API key based)."""

    name = "remote"

    def __init__(
        self,
        base_url: str,
        api_key: str,
        model: str = "text-embedding-3-small",
        dimension: int = EMBEDDING_DIM,
    ) -> None:
        super().__init__(dimension)
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.model = model

    def encode(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        import httpx

        response = httpx.post(
            f"{self.base_url}/embeddings",
            headers={"Authorization": f"Bearer {self.api_key}"},
            json={"model": self.model, "input": texts},
            timeout=30.0,
        )
        response.raise_for_status()
        payload = response.json()
        data = sorted(payload.get("data", []), key=lambda item: item.get("index", 0))
        return [self._normalize([float(x) for x in item["embedding"]]) for item in data]

    @staticmethod
    def _normalize(vec: list[float]) -> list[float]:
        norm = math.sqrt(sum(v * v for v in vec))
        if norm > 0.0:
            return [v / norm for v in vec]
        return vec


_provider: EmbeddingProvider | None = None
_provider_lock = threading.Lock()


def _build_provider() -> EmbeddingProvider:
    """Select the best available provider (respecting ``OBSIGATE_EMBEDDING_PROVIDER``)."""
    requested = os.getenv("OBSIGATE_EMBEDDING_PROVIDER", "auto").strip().lower()

    if requested in ("auto", "local", "sentence-transformers") and SentenceTransformerProvider.is_available():
        model = os.getenv("OBSIGATE_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
        return SentenceTransformerProvider(model)

    remote_key = os.getenv("OBSIGATE_EMBEDDING_API_KEY", "")
    remote_url = os.getenv("OBSIGATE_EMBEDDING_BASE_URL", "")
    if requested in ("auto", "remote") and remote_key and remote_url:
        model = os.getenv("OBSIGATE_EMBEDDING_MODEL", "text-embedding-3-small")
        dim = int(os.getenv("OBSIGATE_EMBEDDING_DIM", str(EMBEDDING_DIM)))
        return RemoteEmbeddingProvider(remote_url, remote_key, model, dim)

    if requested == "remote":
        logger.warning(
            "OBSIGATE_EMBEDDING_PROVIDER=remote but OBSIGATE_EMBEDDING_API_KEY/BASE_URL missing; using hash fallback"
        )
    return HashEmbeddingProvider()


def get_embedding_provider() -> EmbeddingProvider:
    """Return the cached embedding provider (built on first use)."""
    global _provider
    with _provider_lock:
        if _provider is None:
            _provider = _build_provider()
            logger.info("Semantic embedding provider: %s (dim=%d)", _provider.name, _provider.dimension)
        return _provider


def reset_embedding_provider() -> None:
    """Forget the cached provider (used by tests and config reloads)."""
    global _provider
    with _provider_lock:
        _provider = None


# ---------------------------------------------------------------------------
# Vector store
# ---------------------------------------------------------------------------
class VectorStore:
    """In-memory vector store with optional numpy / faiss acceleration.

    Vectors are always kept as Python lists (source of truth). A numpy matrix
    and/or a faiss ``IndexFlatIP`` are built lazily and invalidated on mutation.
    All vectors are expected to be L2-normalized, so the inner product equals
    the cosine similarity.
    """

    def __init__(self, dimension: int = EMBEDDING_DIM) -> None:
        self.dimension = dimension
        self._keys: list[str] = []
        self._chunks: list[str] = []
        self._vectors: list[list[float]] = []
        self._dirty = True
        self._numpy: Any | None = None
        self._numpy_checked = False
        self._matrix: Any | None = None
        self._faiss: Any | None = None
        self._faiss_checked = False
        self._faiss_index: Any | None = None

    def __len__(self) -> int:
        return len(self._vectors)

    def clear(self) -> None:
        self._keys = []
        self._chunks = []
        self._vectors = []
        self._dirty = True

    def add(self, key: str, chunk: str, vector: list[float]) -> None:
        self._keys.append(key)
        self._chunks.append(chunk)
        self._vectors.append(vector)
        self._dirty = True

    def remove_document(self, key: str) -> None:
        """Remove every chunk belonging to *key*."""
        kept = [(k, c, v) for k, c, v in zip(self._keys, self._chunks, self._vectors) if k != key]
        if len(kept) == len(self._vectors):
            return
        self._keys = [k for k, _, _ in kept]
        self._chunks = [c for _, c, _ in kept]
        self._vectors = [v for _, _, v in kept]
        self._dirty = True

    # -- optional accelerators -------------------------------------------------
    def _get_numpy(self) -> Any | None:
        if not self._numpy_checked:
            self._numpy_checked = True
            try:
                import numpy as np

                self._numpy = np
            except Exception:
                self._numpy = None
        return self._numpy

    def _get_faiss(self) -> Any | None:
        if not self._faiss_checked:
            self._faiss_checked = True
            try:
                import faiss

                self._faiss = faiss
            except Exception:
                self._faiss = None
        return self._faiss

    def _rebuild_accelerators(self) -> None:
        self._dirty = False
        np = self._get_numpy()
        if np is None or not self._vectors:
            self._matrix = None
            self._faiss_index = None
            return
        self._matrix = np.asarray(self._vectors, dtype="float32")
        faiss = self._get_faiss()
        if faiss is not None:
            index = faiss.IndexFlatIP(self.dimension)
            index.add(self._matrix)
            self._faiss_index = index
        else:
            self._faiss_index = None

    def search(self, query_vector: list[float], top_k: int = DEFAULT_TOP_K) -> list[tuple[str, float]]:
        """Return ``(doc_key, cosine_similarity)`` pairs sorted by similarity."""
        if not self._vectors:
            return []
        top_k = max(1, min(top_k, len(self._vectors)))
        if self._dirty:
            self._rebuild_accelerators()

        np = self._get_numpy()
        if np is not None and self._faiss_index is not None and self._matrix is not None:
            query = np.asarray([query_vector], dtype="float32")
            scores, indices = self._faiss_index.search(query, top_k)
            return [
                (self._keys[int(idx)], float(score))
                for score, idx in zip(scores[0], indices[0])
                if idx >= 0
            ]

        if np is not None and self._matrix is not None:
            query = np.asarray(query_vector, dtype="float32")
            scores = self._matrix @ query
            order = np.argsort(scores)[::-1][:top_k]
            return [(self._keys[int(i)], float(scores[int(i)])) for i in order]

        scored = [(self._keys[i], _dot(self._vectors[i], query_vector)) for i in range(len(self._vectors))]
        scored.sort(key=lambda item: item[1], reverse=True)
        return scored[:top_k]

    def chunk_of(self, index: int) -> str:
        """Return the stored chunk text at *index* (used by diagnostics/tests)."""
        return self._chunks[index]


def _dot(a: list[float], b: list[float]) -> float:
    """Dot product for two equal-length vectors."""
    return sum(x * y for x, y in zip(a, b))


# ---------------------------------------------------------------------------
# Reciprocal Rank Fusion
# ---------------------------------------------------------------------------
def rrf_fuse(rankings: list[list[str]], k: int = RRF_K) -> dict[str, float]:
    """Fuse several ranked key lists into a single score map.

    ``score(key) = Σ_rankings 1 / (k + rank(key))`` where ``rank`` is
    1-based. Documents ranked highly by several methods rise to the top.

    Args:
        rankings: Ordered lists of document keys (best first).
        k: RRF smoothing constant.

    Returns:
        Mapping ``doc_key -> fused score`` (insertion order is unspecified).
    """
    scores: dict[str, float] = {}
    for ranking in rankings:
        seen: set[str] = set()
        for rank, key in enumerate(ranking):
            if key in seen:
                continue
            seen.add(key)
            scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank + 1)
    return scores


# ---------------------------------------------------------------------------
# Semantic index
# ---------------------------------------------------------------------------
class SemanticIndex:
    """Holds document chunk embeddings and answers similarity queries."""

    def __init__(self, provider: EmbeddingProvider | None = None) -> None:
        self.provider = provider
        self.store = VectorStore(provider.dimension if provider else EMBEDDING_DIM)
        self.doc_keys: set[str] = set()
        self._ready = False
        self._lock = threading.Lock()

    def is_ready(self) -> bool:
        """Return True once a full rebuild has completed."""
        return self._ready

    def _ensure_provider(self) -> EmbeddingProvider:
        if self.provider is None:
            self.provider = get_embedding_provider()
            self.store = VectorStore(self.provider.dimension)
        return self.provider

    @staticmethod
    def _document_text(file_info: dict[str, Any]) -> str:
        title = file_info.get("title", "") or ""
        content = file_info.get("content", "") or ""
        return (title + "\n\n" + content).strip()

    def _embed_document(self, doc_key: str, file_info: dict[str, Any]) -> None:
        text = self._document_text(file_info)
        if not text:
            return
        provider = self._ensure_provider()
        chunks = chunk_text(text)
        if not chunks:
            return
        vectors = provider.encode(chunks)
        for chunk, vector in zip(chunks, vectors):
            self.store.add(doc_key, chunk, vector)
        self.doc_keys.add(doc_key)

    def rebuild(self) -> None:
        """Rebuild the whole index from the global in-memory index."""
        from backend.indexer import index

        provider = self._ensure_provider()
        with self._lock:
            self.store = VectorStore(provider.dimension)
            self.doc_keys = set()
            for vault_name, vault_data in index.items():
                for file_info in vault_data.get("files", []):
                    doc_key = f"{vault_name}::{file_info.get('path', '')}"
                    try:
                        self._embed_document(doc_key, file_info)
                    except Exception as exc:
                        logger.warning("Semantic embedding failed for %s: %s", doc_key, exc)
            self._ready = True
        logger.info(
            "Semantic index built: %d documents, %d chunks (provider=%s)",
            len(self.doc_keys),
            len(self.store),
            provider.name,
        )

    def add_document(self, vault_name: str, path: str, file_info: dict[str, Any]) -> None:
        """Add or refresh a single document (no-op until the index is ready)."""
        if not self._ready or not file_info:
            return
        doc_key = f"{vault_name}::{path}"
        with self._lock:
            self.store.remove_document(doc_key)
            self.doc_keys.discard(doc_key)
            try:
                self._embed_document(doc_key, file_info)
            except Exception as exc:
                logger.warning("Semantic embedding failed for %s: %s", doc_key, exc)

    def remove_document(self, vault_name: str, path: str) -> None:
        """Remove a single document (no-op until the index is ready)."""
        if not self._ready:
            return
        doc_key = f"{vault_name}::{path}"
        with self._lock:
            self.store.remove_document(doc_key)
            self.doc_keys.discard(doc_key)

    def search(
        self,
        query: str,
        vault_filter: str = "all",
        top_k: int = DEFAULT_TOP_K,
    ) -> list[tuple[str, float]]:
        """Return ``(doc_key, best_chunk_similarity)`` pairs, best first."""
        if not self._ready or not query or not query.strip():
            return []
        provider = self._ensure_provider()
        query_vector = provider.encode_one(query[:MAX_QUERY_CHARS])
        hits = self.store.search(query_vector, top_k=max(top_k * 4, top_k))
        best: dict[str, float] = {}
        for doc_key, score in hits:
            if vault_filter != "all" and not doc_key.startswith(vault_filter + "::"):
                continue
            if doc_key not in best or score > best[doc_key]:
                best[doc_key] = score
        ordered = sorted(best.items(), key=lambda item: item[1], reverse=True)
        return ordered[:top_k]


_semantic_index: SemanticIndex | None = None
_index_lock = threading.Lock()


def get_semantic_index() -> SemanticIndex:
    """Return the process-wide semantic index (created on first access)."""
    global _semantic_index
    with _index_lock:
        if _semantic_index is None:
            _semantic_index = SemanticIndex()
        return _semantic_index


def reset_semantic_index() -> None:
    """Drop the singleton index (tests)."""
    global _semantic_index
    with _index_lock:
        _semantic_index = None


def init_semantic_index() -> None:
    """Force a full semantic index build. Called after ``build_index`` on startup."""
    from backend.indexer import index

    if any(vdata.get("files") for vdata in index.values()):
        get_semantic_index().rebuild()


def on_index_change(action: str, vault_name: str, path: str, file_info: dict[str, Any]) -> None:
    """Incremental hook registered with the indexer change notifier."""
    index_obj = get_semantic_index()
    if action == "add" and file_info:
        index_obj.add_document(vault_name, path, file_info)
    elif action == "remove":
        index_obj.remove_document(vault_name, path)


def semantic_search_docs(
    query: str,
    vault_filter: str = "all",
    top_k: int = DEFAULT_TOP_K,
) -> list[tuple[str, float]]:
    """Convenience wrapper around :meth:`SemanticIndex.search`."""
    return get_semantic_index().search(query, vault_filter=vault_filter, top_k=top_k)


def semantic_status() -> dict[str, Any]:
    """Return provider/index diagnostics for the API and the UI."""
    index_obj = get_semantic_index()
    provider = index_obj.provider or get_embedding_provider()
    return {
        "available": index_obj.is_ready(),
        "provider": provider.name,
        "dimension": provider.dimension,
        "documents": len(index_obj.doc_keys),
        "chunks": len(index_obj.store),
    }