"""Optional sqlite-vec semantic-retrieval backend (ticket #000039).

Runs **alongside** the FTS5 routes (never replacing them — see #000039
§4): vec closes the *semantic* allusion gap (query and target chunk
share zero stems but mean the same thing), where the lexical pipeline
closes the *lexical* gap. Hits land ``audit_mode=UNGROUNDED``, same as
FTS5 — vec changes recall, never warrant; embeddings are soft signal
and never enter the proof path (CLAUDE.md soft-hash / hard-hash rule).

Optional dep: ``pip install 'arborist[vec]'`` (sqlite-vec + fastembed).
When ``sqlite_vec`` is not importable, ``VEC_AVAILABLE`` is False and
the CLI must not surface ``--backend vec`` / ``arborist embed``.

Hyperparameters — the "obvious" v1 tuning (#000039 §3, §5):

    embedder        BAAI/bge-small-en-v1.5   (384-dim, CPU, ~130 MB ONNX
                                              via fastembed; ticket §5's
                                              recommended bundled model)
    dimension       384
    quantization    float32                  (1536 bytes/chunk; int8 /
                                              binary are the production
                                              storage knobs per §3.1 —
                                              not wired in v1)
    distance metric cosine                    (bge outputs L2-normalized
                                              vectors, so cosine ranking
                                              ≡ L2 ranking; named cosine
                                              for clarity)
    ANN index       flat                      (vec0 default — exact KNN;
                                              IVF / DiskANN are the
                                              scale escape hatch per §3.1)
    top-k           20                        (matches SearchBackend.search
                                              default ``limit``)

These five fields fold into ``governance_policy_hash`` in a later
phase (like every other versioned default — #000039 §6). The
``VEC_BACKEND_VERSION`` token below is the rotation discriminator;
bumping the embedder / dim / quant / metric / ann re-stales prior
``chunk_vecs`` rows on lookup (the ``vec_meta`` table records which
version populated a shard).
"""

from __future__ import annotations

import sqlite3
from collections.abc import Callable, Iterable, Iterator

from arborist.compress import unpack_chunk
from arborist.search.base import AuditMode, Hit, SearchBackend
from arborist.search.fts5 import _build_snippet

try:  # optional dep guard — keep this cheap, no model load here
    import sqlite_vec  # type: ignore

    VEC_AVAILABLE = True
except ImportError:  # pragma: no cover - exercised in CI matrices w/o the extra
    sqlite_vec = None  # type: ignore
    VEC_AVAILABLE = False


# --- hyperparameters (the v1 "obvious" tuning) -----------------------

EMBED_MODEL = "BAAI/bge-small-en-v1.5"
EMBED_DIM = 384
EMBED_METRIC = "cosine"
ANN_INDEX = "flat"
DEFAULT_TOP_K = 20

# Quantization knob (#000039 §3.1). float32 = +25% corpus tax (the v1
# default — simplest, max fidelity); int8 = ~4× smaller (+6% tax, ~1-3%
# recall hit — the recommended production default). binary is a future
# pass (needs a two-stage re-rank). bge outputs L2-normalized [-1,1]-ish
# components; int8 scales by 127 (theoretical-range scaling — leaves
# headroom but plenty of resolution for cosine ranking).
QUANTS = ("float32", "int8")
EMBED_QUANT = "float32"  # the v1 default
_INT8_SCALE = 127.0

# Per-quant vec0 column element type and the SQL constructor that tags a
# blob as that type. sqlite-vec v0.1.9 treats a bare blob as float32 —
# int8 blobs MUST be wrapped in vec_int8(...); we wrap float32 in
# vec_f32(...) too for symmetry/clarity. (Constructor names come from
# this fixed dict, never from user input — no SQL-injection surface.)
_VEC0_ELEM_TYPE = {"float32": "float", "int8": "int8"}
_VEC_CTOR = {"float32": "vec_f32", "int8": "vec_int8"}


def vec_backend_version(quant: str = EMBED_QUANT) -> str:
    """Rotation discriminator for a given quantization.

    Bumping the model / dim / metric / ann / quant changes this token;
    ``vec_meta.backend_version`` records it per shard so a mismatch is
    detectable at lookup time (and an ``arborist embed --rebuild`` under
    a different quant re-embeds + recreates the table at the new type).
    """
    if quant not in QUANTS:
        raise ValueError(f"quant must be one of {QUANTS}; got {quant!r}")
    return f"vec-v1-bge-small-en-v1.5-{EMBED_DIM}{quant}-{EMBED_METRIC}-{ANN_INDEX}"


# Back-compat constant — the float32 (v1 default) version token.
VEC_BACKEND_VERSION = vec_backend_version(EMBED_QUANT)


# --- extension load + schema -----------------------------------------


def load_vec_extension(conn: sqlite3.Connection) -> None:
    """Load the sqlite-vec loadable extension onto ``conn``.

    ``store.connect()`` does not enable extension loading (it's a
    capability most connections don't need); we toggle it on, load,
    toggle it back off — same pattern as any optional SQLite extension.
    """
    if not VEC_AVAILABLE:
        raise RuntimeError(
            "sqlite-vec is not installed — pip install 'arborist[vec]'"
        )
    conn.enable_load_extension(True)
    try:
        sqlite_vec.load(conn)  # type: ignore[union-attr]
    finally:
        conn.enable_load_extension(False)


def existing_quant(conn: sqlite3.Connection) -> str | None:
    """Return the quant of an existing ``chunk_vecs`` table, or None.

    Parsed from the table's stored CREATE statement (``... embedding
    int8[384] ...`` vs ``... embedding float[384] ...``).
    """
    row = conn.execute(
        "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunk_vecs'"
    ).fetchone()
    if not row or not row[0]:
        return None
    sql = row[0]
    for quant, elem in _VEC0_ELEM_TYPE.items():
        if f"embedding {elem}[" in sql:
            return quant
    return None


def ensure_chunk_vecs_table(conn: sqlite3.Connection, *, quant: str = EMBED_QUANT) -> None:
    """Create the ``chunk_vecs`` vec0 virtual table + ``vec_meta`` if absent.

    ``chunk_vecs`` is a *sibling* table (additive — does not touch
    ``chunks`` / ``documents`` / the audit chain). ``vec_meta`` records
    the ``vec_backend_version(quant)`` that populated this shard so a
    rotation (including a quant change) is detectable.

    Raises ``ValueError`` if ``chunk_vecs`` already exists at a *different*
    quant — switching quant requires recreating the table (use
    ``embed_documents(..., rebuild=True)`` / ``arborist embed --rebuild
    --quant ...``), since the vec0 column element type can't be altered
    in place.
    """
    if quant not in QUANTS:
        raise ValueError(f"quant must be one of {QUANTS}; got {quant!r}")
    load_vec_extension(conn)
    have = existing_quant(conn)
    if have is not None and have != quant:
        raise ValueError(
            f"chunk_vecs already exists at quant={have!r}; cannot switch to "
            f"{quant!r} in place — re-run with --rebuild (DROP + recreate)."
        )
    conn.execute(
        f"CREATE VIRTUAL TABLE IF NOT EXISTS chunk_vecs USING vec0("
        f"  chunk_id INTEGER PRIMARY KEY,"
        f"  embedding {_VEC0_ELEM_TYPE[quant]}[{EMBED_DIM}] distance_metric={EMBED_METRIC}"
        f")"
    )
    conn.execute(
        "CREATE TABLE IF NOT EXISTS vec_meta ("
        "  key TEXT PRIMARY KEY,"
        "  value TEXT NOT NULL"
        ")"
    )
    conn.execute(
        "INSERT OR REPLACE INTO vec_meta(key, value) VALUES ('backend_version', ?)",
        (vec_backend_version(quant),),
    )


# --- embedder --------------------------------------------------------

# Callable: list[str] -> sequence of dim-length float lists (one per text).
Embedder = Callable[[list[str]], "Iterable[Iterable[float]]"]

_DEFAULT_EMBEDDER: Embedder | None = None


def default_embedder() -> Embedder:
    """fastembed-backed bge-small-en-v1.5 embedder, lazily constructed + cached.

    fastembed's ``TextEmbedding.embed`` already yields L2-normalized
    float32 numpy arrays, so the cosine / L2 equivalence holds without
    a re-normalize step. The model (~130 MB ONNX) downloads to the
    fastembed cache on first construction.
    """
    global _DEFAULT_EMBEDDER
    if _DEFAULT_EMBEDDER is not None:
        return _DEFAULT_EMBEDDER
    try:
        from fastembed import TextEmbedding  # type: ignore
    except ImportError as e:  # pragma: no cover
        raise RuntimeError(
            "fastembed is not installed — pip install 'arborist[vec]' "
            "(or pass a custom embedder to VecBackend / embed_documents)"
        ) from e

    model = TextEmbedding(model_name=EMBED_MODEL)

    def _embed(texts: list[str]) -> Iterator[list[float]]:
        for vec in model.embed(texts):
            yield [float(x) for x in vec]

    _DEFAULT_EMBEDDER = _embed
    return _embed


def _to_blob(vec: Iterable[float], quant: str = EMBED_QUANT) -> bytes:
    """Serialize a float vector to sqlite-vec's blob form for ``quant``.

    float32: ``serialize_float32`` verbatim. int8: scale each component
    by 127 (theoretical [-1,1] range), clamp to [-127, 127], round, then
    ``serialize_int8``. The same scaling is applied to the query vector
    so distances are comparable. (Cosine is scale-invariant, so the
    uniform ×127 cancels in the ranking; clamping touches only the rare
    component beyond ±1.)
    """
    vals = list(vec)
    if quant == "float32":
        return sqlite_vec.serialize_float32(vals)  # type: ignore[union-attr]
    if quant == "int8":
        scaled = [
            max(-127, min(127, int(round(v * _INT8_SCALE)))) for v in vals
        ]
        return sqlite_vec.serialize_int8(scaled)  # type: ignore[union-attr]
    raise ValueError(f"quant must be one of {QUANTS}; got {quant!r}")


# --- ingest ----------------------------------------------------------


def embed_documents(
    conn: sqlite3.Connection,
    *,
    embedder: Embedder | None = None,
    limit: int | None = None,
    batch_size: int = 256,
    incremental: bool = True,
    rebuild: bool = False,
    quant: str = EMBED_QUANT,
    progress: Callable[[int, int], None] | None = None,
) -> int:
    """Populate ``chunk_vecs`` for ``chunks`` rows with content.

    Idempotency model (see #000039 — chunk content is immutable per
    ``chunk_id``: a chunk, once embedded, never needs re-embedding
    unless the *embedder* or *quant* changes):

    - ``incremental=True`` (default): embed only chunks **not already
      in chunk_vecs** (``chunk_id NOT IN (SELECT chunk_id FROM chunk_vecs)``).
      This is the after-ingest / cron / Prometheus-Σ-sweep path — it
      picks up exactly the newly-ingested chunks; re-running is cheap
      (no-op once everything's embedded).
    - ``incremental=False``: re-embed every chunk with content
      (delete-then-insert all). The embedder-changed case.
    - ``rebuild=True``: DROP + recreate ``chunk_vecs`` first (at the
      requested ``quant``), then a full pass — the clean
      ``vec_backend_version`` bump *and* the only way to switch quant
      (the vec0 column element type can't be altered in place). A search
      mid-rebuild never mixes configs: it's all new from the recreated
      table, growing as the pass runs. Implies a full (non-incremental)
      pass. **Required when changing ``quant`` on an existing table.**

    ``quant`` ∈ {"float32" (default, +25% corpus tax, max fidelity),
    "int8" (~4× smaller, +6% tax, ~1-3% recall hit — production default)}.

    Cold-evicted chunks (``content IS NULL``) are skipped — the vec
    row, if it exists, stays valid (content is the same on rehydrate).
    ``limit`` caps chunks processed (smoke-test knob). Returns the
    count newly embedded this call.
    """
    if quant not in QUANTS:
        raise ValueError(f"quant must be one of {QUANTS}; got {quant!r}")
    embedder = embedder or default_embedder()
    load_vec_extension(conn)
    if not rebuild and existing_quant(conn) not in (None, quant):
        raise ValueError(
            f"chunk_vecs exists at quant={existing_quant(conn)!r}; pass "
            f"rebuild=True (--rebuild) to switch to {quant!r}."
        )
    if rebuild:
        try:
            conn.execute("DROP TABLE IF EXISTS chunk_vecs")
        except sqlite3.OperationalError:
            pass
        incremental = False  # recreated table is empty → full pass anyway
    ensure_chunk_vecs_table(conn, quant=quant)

    where = "content IS NOT NULL"
    if incremental:
        where += " AND chunk_id NOT IN (SELECT chunk_id FROM chunk_vecs)"

    total_known = conn.execute(
        f"SELECT COUNT(*) FROM chunks WHERE {where}"
    ).fetchone()[0]
    if limit is not None:
        total_known = min(total_known, limit)

    done = 0
    batch_ids: list[int] = []
    batch_texts: list[str] = []

    def flush() -> None:
        nonlocal done
        if not batch_texts:
            return
        vecs = list(embedder(batch_texts))
        conn.execute("BEGIN")
        try:
            # vec0 virtual tables don't honor INSERT OR REPLACE — re-inserting
            # an existing chunk_id is a hard UNIQUE error. Delete-then-insert
            # is the upsert; harmless no-op on rows that aren't present yet
            # (the incremental path) and the re-embed path otherwise.
            conn.executemany(
                "DELETE FROM chunk_vecs WHERE chunk_id = ?",
                [(cid,) for cid in batch_ids],
            )
            conn.executemany(
                f"INSERT INTO chunk_vecs(chunk_id, embedding) "
                f"VALUES (?, {_VEC_CTOR[quant]}(?))",
                [(cid, _to_blob(v, quant)) for cid, v in zip(batch_ids, vecs)],
            )
            conn.execute("COMMIT")
        except Exception:
            conn.execute("ROLLBACK")
            raise
        done += len(batch_ids)
        if progress is not None:
            progress(done, total_known)
        batch_ids.clear()
        batch_texts.clear()

    rows = conn.execute(
        f"SELECT chunk_id, content FROM chunks WHERE {where} ORDER BY chunk_id"
        + (f" LIMIT {int(limit)}" if limit is not None else "")
    )
    for row in rows:
        text = unpack_chunk(row["content"])
        if not text:
            continue
        batch_ids.append(int(row["chunk_id"]))
        batch_texts.append(text)
        if len(batch_texts) >= batch_size:
            flush()
    flush()
    return done


# --- search backend --------------------------------------------------


class VecBackend(SearchBackend):
    """ANN-over-``chunk_vecs`` retrieval. UNGROUNDED hits, same as FTS5."""

    name = "vec"
    audit_mode = AuditMode.UNGROUNDED

    def __init__(
        self,
        conn: sqlite3.Connection,
        embedder: Embedder | None = None,
        *,
        quant: str | None = None,
    ):
        super().__init__(conn)
        load_vec_extension(conn)
        self._embedder = embedder  # lazy: default constructed on first search
        # The quant the existing chunk_vecs was populated at; falls back
        # to the v1 default if the table doesn't exist yet (search will
        # return [] in that case anyway — see populated()).
        self.quant = quant or existing_quant(conn) or EMBED_QUANT

    @property
    def embedder(self) -> Embedder:
        if self._embedder is None:
            self._embedder = default_embedder()
        return self._embedder

    def populated(self) -> bool:
        """True if ``chunk_vecs`` exists and has at least one row."""
        try:
            n = self.conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0]
        except sqlite3.OperationalError:
            return False
        return n > 0

    def search(self, query: str, limit: int = DEFAULT_TOP_K) -> list[Hit]:
        if not query.strip():
            return []
        if not self.populated():
            return []
        qvec = next(iter(self.embedder([query])))
        rows = self.conn.execute(
            f"""
            SELECT
                cv.chunk_id,
                cv.distance        AS distance,
                c.document_root,
                c.idx,
                c.content          AS raw_content,
                d.document_uri,
                d.title
            FROM chunk_vecs AS cv
            JOIN chunks    AS c ON c.chunk_id = cv.chunk_id
            JOIN documents AS d ON d.document_root = c.document_root
            WHERE cv.embedding MATCH {_VEC_CTOR[self.quant]}(?)
              AND k = ?
            ORDER BY cv.distance
            """,
            (_to_blob(qvec, self.quant), int(limit)),
        ).fetchall()
        return [
            Hit(
                document_root=r["document_root"],
                document_uri=r["document_uri"],
                chunk_idx=r["idx"],
                snippet=_build_snippet(unpack_chunk(r["raw_content"]) or "", query),
                # cosine/L2 distance: lower = better. Flip sign so the
                # merge layer (which treats higher score as better, per
                # FTS5Backend) ranks vec hits correctly. Bounded; not
                # comparable to BM25 scale — hybrid fusion uses RRF
                # (#000039 §4.2), not raw-score addition.
                score=-float(r["distance"]) if r["distance"] is not None else 0.0,
                audit_mode=self.audit_mode,
                title=r["title"],
            )
            for r in rows
        ]


__all__ = [
    "VEC_AVAILABLE",
    "VEC_BACKEND_VERSION",
    "EMBED_MODEL",
    "EMBED_DIM",
    "EMBED_QUANT",
    "QUANTS",
    "vec_backend_version",
    "existing_quant",
    "VecBackend",
    "embed_documents",
    "ensure_chunk_vecs_table",
    "load_vec_extension",
    "default_embedder",
]
