"""
BIBIVIDI — Matcher
Embeds incoming prompts and finds the best matching script
using cosine similarity against the library's embeddings.

Index backend
─────────────
  faiss installed → FAISS IndexFlatIP on L2-normalised vectors
                    (inner product of unit vectors = cosine similarity).
                    When the library grows to ≥ FAISS_IVF_THRESHOLD scripts,
                    the index is upgraded to IndexIVFFlat (16 centroids) for
                    sub-linear approximate nearest-neighbour search.
  faiss absent    → O(n) numpy cosine scan (original behaviour, always correct).

Install the FAISS extra:  pip install "bibividi[faiss]"
"""
from __future__ import annotations

import asyncio
import json as _json
import logging
import re as _re
from collections import deque
from dataclasses import dataclass

import numpy as np

from bvd.config import settings
from bvd.engine.library import Script, library

logger = logging.getLogger(__name__)

# Switch from exact flat index to IVF (approximate) above this many scripts.
# IVF requires at least nlist training vectors; 4× nlist is the minimum
# recommended by FAISS docs.  nlist=16 → threshold=64.
_FAISS_IVF_THRESHOLD = 64
_FAISS_NLIST = 16
_EMBED_DIM = 384  # all-MiniLM-L6-v2 output dimension

try:
    import faiss as _faiss  # type: ignore[import-untyped]
    _FAISS_AVAILABLE = True
    logger.debug("FAISS available — using accelerated ANN index")
except ImportError:
    _faiss = None  # type: ignore[assignment]
    _FAISS_AVAILABLE = False
    logger.debug("FAISS not installed — using numpy cosine scan (pip install 'bibividi[faiss]')")


@dataclass
class MatchResult:
    script: Script
    script_id: str
    score: float


@dataclass
class AssemblyResult:
    """Result of BSS 2.0 auto-assembly: a type-compatible pipeline."""
    pipeline: list[Script]   # ordered list of scripts to execute
    input_type: str          # detected type of the incoming prompt
    output_type: str         # always "plain_text" for the current engine


class Matcher:
    """
    Loads sentence-transformers once, keeps embeddings in memory.
    Refreshes the index from DB on demand (after invalidate() or first use).
    """

    def __init__(self) -> None:
        self._model = None          # lazy-loaded
        # numpy fallback store
        self._index: list[tuple[str, np.ndarray, Script]] = []
        # FAISS store — parallel to _index when FAISS is available
        self._faiss_index: object | None = None   # faiss.Index
        self._index_dirty = True

    # ── Public API ────────────────────────────────────────────

    async def match(self, prompt: str) -> MatchResult | None:
        """Return the best matching script, or None if below threshold."""
        if not self._index or self._index_dirty:
            await self._rebuild_index()

        if not self._index:
            return None

        model = self._get_model()
        prompt_vec = await asyncio.get_event_loop().run_in_executor(
            None, model.encode, [prompt]
        )
        prompt_vec = prompt_vec[0].astype(np.float32)  # shape (dim,)

        if _FAISS_AVAILABLE and self._faiss_index is not None:
            return self._match_faiss(prompt_vec)

        return self._match_numpy(prompt_vec)

    async def embed_script(self, script: Script) -> Script:
        """Compute and attach an embedding to a script (call before saving)."""
        model = self._get_model()
        tag_parts = [t[len("trigger:"):] if t.startswith("trigger:") else t
                     for t in script.tags]
        tag_str = ". ".join(tag_parts)
        text = f"{script.name}. {script.description}. {tag_str}".strip(". ")
        vec = await asyncio.get_event_loop().run_in_executor(
            None, model.encode, [text]
        )
        script.embedding = vec[0].astype(np.float32).tobytes()
        self._index_dirty = True
        return script

    def invalidate(self) -> None:
        """Force index rebuild on next match call."""
        self._index_dirty = True

    async def assemble_pipeline(self, prompt: str) -> AssemblyResult | None:
        """BSS 2.0: auto-assemble a type-compatible skill pipeline for this prompt.

        Detects the prompt's content type, then BFS-searches installed typed
        skills to find the shortest transformation path to ``"plain_text"``.

        Returns ``AssemblyResult`` with the ordered pipeline, or ``None`` if:
        - The prompt is already plain_text
        - No typed skills are installed
        - No type-compatible path exists within 5 hops
        """
        # Cheap type check first — avoids rebuilding the index for plain text
        input_type = _detect_input_type(prompt)
        if input_type == "plain_text":
            return None

        if not self._index or self._index_dirty:
            await self._rebuild_index()

        typed_scripts = [s for _, _, s in self._index if s.inputs_type and s.outputs_type]
        if not typed_scripts:
            return None

        from bvd.engine.library import Library
        graph = Library.build_type_graph(typed_scripts)

        path = _bfs_type_path(graph, input_type, "plain_text")
        if not path:
            return None

        logger.debug(
            "Auto-assembly found %d-skill path: %s → plain_text (%s)",
            len(path),
            input_type,
            " → ".join(s.id for s in path),
        )
        return AssemblyResult(pipeline=path, input_type=input_type, output_type="plain_text")

    # ── Internal ──────────────────────────────────────────────

    def _get_model(self):
        if self._model is None:
            logger.info("Loading embedding model: %s", settings.embedding_model)
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(settings.embedding_model)
            logger.info("Embedding model loaded")
        return self._model

    async def _rebuild_index(self) -> None:
        scripts = await library.all_with_embeddings()
        self._index = []
        for s in scripts:
            if s.embedding:
                vec = np.frombuffer(s.embedding, dtype=np.float32).copy()
                self._index.append((s.id, vec, s))

        if _FAISS_AVAILABLE and self._index:
            self._faiss_index = _build_faiss_index(self._index)
            logger.debug(
                "FAISS index built: %d vectors (%s)",
                len(self._index),
                "IVF" if len(self._index) >= _FAISS_IVF_THRESHOLD else "Flat",
            )
        else:
            self._faiss_index = None

        self._index_dirty = False
        logger.debug("Matcher index rebuilt: %d scripts", len(self._index))

    def _match_faiss(self, prompt_vec: np.ndarray) -> MatchResult | None:
        """Search using the FAISS index (cosine via normalised inner product)."""
        assert _faiss is not None
        query = _l2_normalise(prompt_vec).reshape(1, -1)
        scores, ids = self._faiss_index.search(query, 1)  # type: ignore[union-attr]
        if ids[0][0] == -1:
            return None
        best_score = float(scores[0][0])
        if best_score < settings.match_threshold:
            return None
        script_id, _, script = self._index[ids[0][0]]
        return MatchResult(script=script, script_id=script_id, score=best_score)

    def _match_numpy(self, prompt_vec: np.ndarray) -> MatchResult | None:
        """O(n) cosine scan — used when FAISS is not installed."""
        best_score = -1.0
        best_entry = None
        for script_id, emb, script in self._index:
            score = float(_cosine(prompt_vec, emb))
            if score > best_score:
                best_score = score
                best_entry = (script_id, script)

        if best_score >= settings.match_threshold and best_entry:
            script_id, script = best_entry
            return MatchResult(script=script, script_id=script_id, score=best_score)
        return None


# ── FAISS helpers ─────────────────────────────────────────────────────────────

def _l2_normalise(v: np.ndarray) -> np.ndarray:
    norm = np.linalg.norm(v)
    return v / norm if norm > 0 else v


def _build_faiss_index(
    index: list[tuple[str, np.ndarray, Script]],
) -> object:  # faiss.Index
    """Build a FAISS index from the in-memory script list.

    - n <  _FAISS_IVF_THRESHOLD → IndexFlatIP (exact, BLAS-accelerated)
    - n >= _FAISS_IVF_THRESHOLD → IndexIVFFlat (approximate, 16 centroids)

    All vectors are L2-normalised so inner product == cosine similarity.
    """
    assert _faiss is not None
    vecs = np.stack([_l2_normalise(emb) for _, emb, _ in index]).astype(np.float32)
    n, dim = vecs.shape

    if n < _FAISS_IVF_THRESHOLD:
        idx = _faiss.IndexFlatIP(dim)
        idx.add(vecs)
        return idx

    # IVF: train with all available vectors, search 4 of nlist cells
    quantiser = _faiss.IndexFlatIP(dim)
    idx = _faiss.IndexIVFFlat(quantiser, dim, _FAISS_NLIST, _faiss.METRIC_INNER_PRODUCT)
    idx.train(vecs)
    idx.add(vecs)
    idx.nprobe = 4  # cells visited per query — trade recall for speed
    return idx


# ── BSS 2.0 auto-assembly helpers ─────────────────────────────────────────────

def _detect_input_type(prompt: str) -> str:
    """Heuristic content-type detection from prompt text.

    Returns one of: "pdf", "diff", "csv", "json", "xml", "plain_text".
    Operates on the first 1 000 chars to keep latency negligible.
    """
    head = prompt[:1000]

    # PDF: binary magic bytes or base64-encoded PDF ("JVBERi0" = "%PDF-" in base64)
    if "%PDF-" in head or head.lstrip().startswith("JVBERi0"):
        return "pdf"

    # Unified diff: presence of diff headers
    if _re.search(r"^(---|\+\+\+|@@\s)", head, _re.MULTILINE):
        return "diff"

    # JSON: valid top-level object or array
    stripped = head.lstrip()
    if stripped.startswith(("{", "[")):
        try:
            _json.loads(prompt[:4000])
            return "json"
        except (ValueError, _json.JSONDecodeError):
            pass

    # XML
    if stripped.startswith("<?xml") or _re.match(r"<[A-Za-z]", stripped):
        return "xml"

    # CSV: ≥3 of the first 5 non-empty lines have ≥2 commas
    lines = [ln for ln in head.splitlines() if ln.strip()]
    if len(lines) >= 3 and sum(1 for ln in lines[:5] if ln.count(",") >= 2) >= 3:
        return "csv"

    return "plain_text"


def _bfs_type_path(
    graph: dict[str, list[tuple[str, Script]]],
    start: str,
    target: str,
    max_depth: int = 5,
) -> list[Script] | None:
    """BFS shortest path through the type-compatibility graph.

    Each node is a content-type string; each edge is a ``Script`` that
    transforms ``inputs_type → outputs_type``.

    Returns the ordered list of ``Script`` objects forming the shortest path
    from ``start`` to ``target``, or ``None`` if no path exists within
    ``max_depth`` steps.
    """
    if start == target:
        return []

    # State: (current_type, accumulated_path)
    queue: deque[tuple[str, list[Script]]] = deque([(start, [])])
    visited: set[str] = {start}

    while queue:
        current_type, path = queue.popleft()
        if len(path) >= max_depth:
            continue

        for output_type, skill in graph.get(current_type, []):
            new_path = path + [skill]
            if output_type == target:
                return new_path
            if output_type not in visited:
                visited.add(output_type)
                queue.append((output_type, new_path))

    return None


# ── numpy cosine (kept for fallback + tests) ──────────────────────────────────

def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    denom = (np.linalg.norm(a) * np.linalg.norm(b))
    if denom == 0:
        return 0.0
    return float(np.dot(a, b) / denom)


# Singleton
matcher = Matcher()
