# utils/sqlite_vec_index.py — sqlite-vec vec0 ANN mirror for session_search_db.vector_chunks.
from __future__ import annotations

import logging
import sqlite3
from typing import Any

_log = logging.getLogger(__name__)

_VEC_TABLE = "vec_chunks"
_PROBE_LOCK = __import__("threading").Lock()
_extension_installed: bool | None = None
_vec_backend: str = ""
_distance_metric: str = "cosine"
_ensured_dim: int | None = None

_BACKFILL_CHUNK = 500


def probe_backend_string() -> str:
    """Raw probe version string (e.g. sqlite_vec:v0.1.9)."""
    return _vec_backend


def vec_backend_label() -> str:
    """Human-readable backend label for operator status APIs."""
    if not _extension_installed:
        return "bruteforce"
    metric = _distance_metric or "cosine"
    ver = _vec_backend or "sqlite_vec"
    return f"vec0:{metric} ({ver})"


def distance_metric() -> str:
    return _distance_metric


def _probe_extension_installed() -> bool:
    """Return whether the sqlite-vec extension is present (always probed on :memory:)."""
    global _extension_installed, _vec_backend
    with _PROBE_LOCK:
        if _extension_installed is not None:
            return _extension_installed
        conn = sqlite3.connect(":memory:")
        try:
            conn.enable_load_extension(True)
            import sqlite_vec

            sqlite_vec.load(conn)
            row = conn.execute("SELECT vec_version()").fetchone()
            ver = str(row[0] if row else "unknown")
            _vec_backend = f"sqlite_vec:{ver}"
            _extension_installed = True
        except Exception as exc:
            _log.debug("sqlite-vec extension unavailable: %s", exc)
            _vec_backend = ""
            _extension_installed = False
        finally:
            try:
                conn.enable_load_extension(False)
            except Exception:
                pass
            conn.close()
        return bool(_extension_installed)


def vec_available(conn: sqlite3.Connection | None = None) -> bool:
    """True when the extension is installed and (if *conn* given) loadable on that connection."""
    if not _probe_extension_installed():
        return False
    if conn is None:
        return True
    return load_into(conn)


def _vec_loaded_on_conn(conn: sqlite3.Connection) -> bool:
    try:
        conn.execute("SELECT vec_version()")
        return True
    except Exception:
        return False


def load_into(conn: sqlite3.Connection) -> bool:
    """Load sqlite-vec on *conn* when the extension is installed."""
    if not _probe_extension_installed():
        return False
    if _vec_loaded_on_conn(conn):
        return True
    try:
        conn.enable_load_extension(True)
        import sqlite_vec

        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
        return _vec_loaded_on_conn(conn)
    except Exception as exc:
        _log.debug("sqlite-vec load_into failed: %s", exc)
        try:
            conn.enable_load_extension(False)
        except Exception:
            pass
        return _vec_loaded_on_conn(conn)


def _table_exists(conn: sqlite3.Connection) -> bool:
    row = conn.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
        (_VEC_TABLE,),
    ).fetchone()
    return row is not None


def reset_table_cache() -> None:
    """Clear per-process vec0 DDL cache (call when opening a different search DB file)."""
    global _ensured_dim
    _ensured_dim = None


def _table_ddl(conn: sqlite3.Connection) -> str:
    if not _table_exists(conn):
        return ""
    try:
        row = conn.execute(
            "SELECT sql FROM sqlite_master WHERE type='table' AND name=?",
            (_VEC_TABLE,),
        ).fetchone()
        return str(row[0] if row else "")
    except Exception:
        return ""


def sync_distance_metric_from_ddl(conn: sqlite3.Connection) -> None:
    global _distance_metric
    ddl = _table_ddl(conn)
    if not ddl:
        return
    if "distance_metric=cosine" in ddl:
        _distance_metric = "cosine"
    elif "float[" in ddl:
        _distance_metric = "l2"


def _table_dim_matches(conn: sqlite3.Connection, dim_i: int) -> bool:
    ddl = _table_ddl(conn)
    return f"float[{dim_i}]" in ddl if ddl else False


def ensure_vec_table(conn: sqlite3.Connection, dim: int) -> bool:
    """Create vec0 mirror table at *dim* if missing; record metric in module state."""
    global _distance_metric, _ensured_dim
    dim_i = int(dim)
    if dim_i < 1:
        return False
    if not load_into(conn):
        return False
    if _table_exists(conn) and _table_dim_matches(conn, dim_i):
        _ensured_dim = dim_i
        sync_distance_metric_from_ddl(conn)
        return True
    if _table_exists(conn) and not _table_dim_matches(conn, dim_i):
        try:
            conn.execute(f"DROP TABLE IF EXISTS {_VEC_TABLE}")
        except Exception:
            pass
        _ensured_dim = None
    for metric in ("cosine", None):
        try:
            if metric:
                sql = (
                    f"CREATE VIRTUAL TABLE IF NOT EXISTS {_VEC_TABLE} "
                    f"USING vec0(embedding float[{dim_i}] distance_metric={metric}, thread_id text)"
                )
            else:
                sql = (
                    f"CREATE VIRTUAL TABLE IF NOT EXISTS {_VEC_TABLE} "
                    f"USING vec0(embedding float[{dim_i}], thread_id text)"
                )
            conn.execute(sql)
            _distance_metric = metric or "l2"
            _ensured_dim = dim_i
            return True
        except Exception as exc:
            _log.debug("ensure_vec_table metric=%s failed: %s", metric, exc)
            continue
    return False


def vec_row_count(conn: sqlite3.Connection) -> int:
    if not _table_exists(conn):
        return 0
    try:
        row = conn.execute(f"SELECT COUNT(*) FROM {_VEC_TABLE}").fetchone()
        return int(row[0] if row else 0)
    except Exception:
        return 0


def vec_thread_row_count(conn: sqlite3.Connection, thread_id: str) -> int:
    tid = (thread_id or "").strip()
    if not tid or not _table_exists(conn):
        return 0
    try:
        row = conn.execute(
            f"SELECT COUNT(*) FROM {_VEC_TABLE} WHERE thread_id = ?",
            (tid,),
        ).fetchone()
        return int(row[0] if row else 0)
    except Exception:
        return 0


def _mirrorable_embedding_sql(
    expected_dim: int | None,
    *,
    table_alias: str = "vc",
) -> tuple[str, list[Any]]:
    """SQL fragment + params: only rows we can mirror into vec0."""
    col = f"{table_alias}.embedding" if (table_alias or "").strip() else "embedding"
    if expected_dim is not None and int(expected_dim) > 0:
        need = int(expected_dim) * 4
        return f" AND {col} IS NOT NULL AND length({col}) = ?", [need]
    return f" AND {col} IS NOT NULL AND length({col}) > 0", []


def purge_orphan_vec_mirrors(conn: sqlite3.Connection) -> int:
    """Remove vec_chunks rows with no matching vector_chunks rowid (stale ANN slots)."""
    if not _table_exists(conn):
        return 0
    try:
        cur = conn.execute(
            f"""
            DELETE FROM {_VEC_TABLE}
            WHERE rowid NOT IN (SELECT rowid FROM vector_chunks)
            """
        )
        return int(cur.rowcount or 0)
    except Exception as exc:
        _log.debug("purge_orphan_vec_mirrors failed: %s", exc)
        return 0


def embedding_blob_mismatch_count(
    conn: sqlite3.Connection,
    expected_dim: int,
    thread_id: str | None = None,
) -> int:
    """Rows whose embedding byte length does not match *expected_dim* (ANN must stay off)."""
    dim_i = int(expected_dim)
    if dim_i < 1:
        return 0
    need = dim_i * 4
    try:
        if thread_id:
            tid = (thread_id or "").strip()
            row = conn.execute(
                """
                SELECT COUNT(*) FROM vector_chunks
                WHERE thread_id = ?
                  AND embedding IS NOT NULL
                  AND length(embedding) > 0
                  AND length(embedding) != ?
                """,
                (tid, need),
            ).fetchone()
        else:
            row = conn.execute(
                """
                SELECT COUNT(*) FROM vector_chunks
                WHERE embedding IS NOT NULL
                  AND length(embedding) > 0
                  AND length(embedding) != ?
                """,
                (need,),
            ).fetchone()
        return int(row[0] if row else 0)
    except Exception:
        return 0


def unmirrored_row_count(
    conn: sqlite3.Connection,
    thread_id: str | None = None,
    *,
    expected_dim: int | None = None,
) -> int:
    """Rows in vector_chunks with no vec_chunks mirror (source of truth for drift)."""
    if not _table_exists(conn):
        emb_sql, emb_params = _mirrorable_embedding_sql(expected_dim, table_alias="")
        try:
            if thread_id:
                row = conn.execute(
                    "SELECT COUNT(*) FROM vector_chunks WHERE thread_id = ?" + emb_sql,
                    [(thread_id or "").strip(), *emb_params],
                ).fetchone()
            else:
                row = conn.execute(
                    "SELECT COUNT(*) FROM vector_chunks WHERE 1=1" + emb_sql,
                    emb_params,
                ).fetchone()
            return int(row[0] if row else 0)
        except Exception:
            return 0
    emb_sql, emb_params = _mirrorable_embedding_sql(expected_dim)
    try:
        if thread_id:
            tid = (thread_id or "").strip()
            row = conn.execute(
                f"""
                SELECT COUNT(*) FROM vector_chunks vc
                WHERE vc.thread_id = ?
                  AND vc.rowid NOT IN (SELECT rowid FROM {_VEC_TABLE})
                """
                + emb_sql,
                (tid, *emb_params),
            ).fetchone()
        else:
            row = conn.execute(
                f"""
                SELECT COUNT(*) FROM vector_chunks vc
                WHERE vc.rowid NOT IN (SELECT rowid FROM {_VEC_TABLE})
                """
                + emb_sql,
                emb_params,
            ).fetchone()
        return int(row[0] if row else 0)
    except Exception:
        return 0


def mirror_insert(
    conn: sqlite3.Connection,
    rowid: int,
    thread_id: str,
    blob: bytes,
    *,
    expected_dim: int | None = None,
) -> bool:
    """Return True when a row was written to vec_chunks."""
    if not _table_exists(conn) or not load_into(conn):
        return False
    if expected_dim is not None and expected_dim > 0:
        expected_bytes = int(expected_dim) * 4
        if len(blob) != expected_bytes:
            _log.warning(
                "[search] vec mirror_insert skipped rowid=%s: blob len %d != %d",
                rowid,
                len(blob),
                expected_bytes,
            )
            return False
    conn.execute(f"DELETE FROM {_VEC_TABLE} WHERE rowid = ?", (int(rowid),))
    conn.execute(
        f"INSERT INTO {_VEC_TABLE}(rowid, embedding, thread_id) VALUES (?, ?, ?)",
        (int(rowid), blob, thread_id),
    )
    return True


def mirror_delete(conn: sqlite3.Connection, rowid: int) -> None:
    if not _table_exists(conn):
        return
    conn.execute(f"DELETE FROM {_VEC_TABLE} WHERE rowid = ?", (int(rowid),))


def mirror_delete_rowids(conn: sqlite3.Connection, rowids: list[int]) -> None:
    if not rowids or not _table_exists(conn):
        return
    placeholders = ",".join("?" * len(rowids))
    conn.execute(
        f"DELETE FROM {_VEC_TABLE} WHERE rowid IN ({placeholders})",
        [int(r) for r in rowids],
    )


def mirror_delete_thread(conn: sqlite3.Connection, thread_id: str) -> None:
    tid = (thread_id or "").strip()
    if not tid or not _table_exists(conn):
        return
    rowids = [
        int(r[0])
        for r in conn.execute(
            "SELECT rowid FROM vector_chunks WHERE thread_id = ?",
            (tid,),
        ).fetchall()
    ]
    mirror_delete_rowids(conn, rowids)


def knn(
    conn: sqlite3.Connection,
    qblob: bytes,
    *,
    k: int,
    thread_id: str | None = None,
) -> list[tuple[int, float]]:
    if not _table_exists(conn) or k < 1:
        return []
    k = max(1, min(int(k), 50000))
    try:
        if thread_id:
            tid = (thread_id or "").strip()
            rows = conn.execute(
                f"""
                SELECT rowid, distance FROM {_VEC_TABLE}
                WHERE embedding MATCH ? AND k = ? AND thread_id = ?
                """,
                (qblob, k, tid),
            ).fetchall()
        else:
            rows = conn.execute(
                f"""
                SELECT rowid, distance FROM {_VEC_TABLE}
                WHERE embedding MATCH ? AND k = ?
                """,
                (qblob, k),
            ).fetchall()
    except Exception as exc:
        _log.debug("vec0 KNN query failed: %s", exc)
        return []
    out: list[tuple[int, float]] = []
    for rowid, dist in rows:
        try:
            out.append((int(rowid), float(dist)))
        except (TypeError, ValueError):
            continue
    return out


def distance_to_semantic_score(distance: float) -> float:
    """Map vec0 distance to cosine-like similarity in [0, 1]."""
    d = float(distance)
    if _distance_metric == "cosine":
        return max(0.0, 1.0 - d)
    # L2 on unit vectors: sim = 1 - d^2/2
    return max(0.0, 1.0 - (d * d) / 2.0)


def backfill(
    conn: sqlite3.Connection,
    *,
    chunk_size: int = _BACKFILL_CHUNK,
    expected_dim: int | None = None,
) -> int:
    """Mirror vector_chunks rows missing from vec_chunks; idempotent."""
    if not _table_exists(conn):
        return 0
    if expected_dim is not None and embedding_blob_mismatch_count(conn, int(expected_dim)) > 0:
        return 0
    purge_orphan_vec_mirrors(conn)
    emb_sql, emb_params = _mirrorable_embedding_sql(expected_dim)
    total = 0
    while True:
        rows = conn.execute(
            f"""
            SELECT vc.rowid, vc.embedding, vc.thread_id
            FROM vector_chunks vc
            WHERE vc.rowid NOT IN (SELECT rowid FROM {_VEC_TABLE})
            """
            + emb_sql
            + """
            LIMIT ?
            """,
            (*emb_params, max(1, int(chunk_size))),
        ).fetchall()
        if not rows:
            break
        progress = 0
        for rowid, blob, tid in rows:
            if not blob:
                continue
            try:
                if mirror_insert(
                    conn,
                    int(rowid),
                    str(tid or ""),
                    blob,
                    expected_dim=expected_dim,
                ):
                    progress += 1
                    total += 1
            except Exception as exc:
                _log.warning(
                    "[search] vec backfill rowid=%s failed: %s",
                    rowid,
                    str(exc)[:120],
                )
        if progress == 0:
            break
    return total


def reset_probe_cache_for_tests() -> None:
    """Test-only: clear cached probe result."""
    global _extension_installed, _vec_backend, _ensured_dim, _distance_metric
    with _PROBE_LOCK:
        _extension_installed = None
        _vec_backend = ""
        _ensured_dim = None
        _distance_metric = "cosine"
