"""LanceDB embedded vector backend — IVF_PQ compatibility path, no daemon.

Replaces numpy brute-force for vector search when --extra hybrid installed.
LanceDB stores vectors in Lance columnar format (memory-mapped, fast load).
Embedded — no server process, no daemon, no network.

Architecture:
  SQLite FTS5 (BM25) + LanceDB (vector IVF_PQ when indexed) = hybrid search.
  Exact NumPy remains the default generation backend; Lance is compatibility.

Install: uv sync --extra hybrid
Storage: cache/lancedb/ (gitignored, rebuildable from Markdown)
"""
from __future__ import annotations

import re
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))
from memory_state import STATE_ROOT  # noqa: E402

LANCEDB_DIR = STATE_ROOT / "cache" / "lancedb"
TABLE_NAME = "pages_vec"
STAGING_TABLE_NAME = "pages_vec_staging"
EMBEDDING_DIM = 384
EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"
EMBEDDING_MODEL_REVISION = "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a"

_lancedb: object | None = None  # lazy connection
_PROJECT_SAFE = re.compile(r"^[A-Za-z0-9_.-]+$")


def _have_lancedb() -> bool:
    """Check if lancedb is importable."""
    try:
        import lancedb  # noqa: F401

        return True
    except ImportError:
        return False


def _get_db():
    """Get or create LanceDB connection. Returns None if unavailable."""
    global _lancedb
    if _lancedb is not None:
        return _lancedb
    if not _have_lancedb():
        return None
    try:
        import lancedb

        LANCEDB_DIR.mkdir(parents=True, exist_ok=True)
        _lancedb = lancedb.connect(str(LANCEDB_DIR))
        return _lancedb
    except Exception:
        return None


def have_lancedb() -> bool:
    """Quick probe: is LanceDB available and has data?"""
    db = _get_db()
    if db is None:
        return False
    try:
        tables = db.table_names()
        return TABLE_NAME in tables
    except Exception:
        return False


def distance_to_similarity(distance: float) -> float:
    """Convert Lance smaller-is-better distance to larger-is-better similarity."""
    if distance < 0:
        distance = 0.0
    return 1.0 / (1.0 + float(distance))


def _rows_from_lance_hits(raw_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Normalize Lance hits: keep lance_distance raw; score is similarity."""
    rows: list[dict[str, Any]] = []
    for item in raw_rows:
        distance = float(item.get("_distance", item.get("lance_distance", 1.0)))
        similarity = round(distance_to_similarity(distance), 6)
        rows.append(
            {
                "path": item.get("path", ""),
                "title": item.get("title", ""),
                "summary": (item.get("summary") or "")[:120],
                "project": item.get("project", "") or "",
                "timestamp": item.get("timestamp", "") or "",
                "status": item.get("status", "") or "",
                "valid_from": item.get("valid_from", "") or "",
                "valid_to": item.get("valid_to", "") or "",
                "authority": item.get("authority", "") or "",
                "candidate_id": item.get("candidate_id")
                or item.get("chunk_id")
                or Path(str(item.get("path", ""))).stem,
                "source_sha256": item.get("source_sha256", "") or "",
                "model": item.get("model", "") or "",
                "model_revision": item.get("model_revision", "") or "",
                "lance_distance": round(distance, 6),
                "vector_score": similarity,
                "score": similarity,
            }
        )
    rows.sort(key=lambda row: (-float(row["vector_score"]), str(row["path"])))
    return rows


def apply_vector_filters(
    rows: list[dict[str, Any]],
    *,
    project: str | None = None,
    since: str | None = None,
    as_of: str | None = None,
    scope: str = "all",
    authority: str | None = None,
) -> list[dict[str, Any]]:
    """Apply the shared hard-filter contract (parity with lexical/NumPy)."""
    from search_memory import apply_hard_filters

    return apply_hard_filters(
        rows,
        project=project,
        since=since,
        as_of=as_of,
        scope=scope,
        authority=authority,
    )


def _sql_literal(value: str) -> str:
    return "'" + value.replace("'", "''") + "'"


def _lance_filter(
    *,
    project: str | None,
    since: str | None,
    as_of: str | None,
    scope: str,
    authority: str | None,
    model_id: str,
    model_revision: str,
) -> str:
    clauses = [
        f"model = {_sql_literal(model_id)}",
        f"model_revision = {_sql_literal(model_revision)}",
    ]
    if project:
        clauses.append(f"lower(project) = lower({_sql_literal(project)})")
    if scope in {"wiki", "memory", "knowledge"}:
        clauses.append("path LIKE 'knowledge/notes/%'")
    if authority:
        clauses.append(f"lower(authority) = lower({_sql_literal(authority)})")
    if as_of:
        date = _sql_literal(as_of[:10])
        clauses.extend(
            (
                f"(timestamp IS NULL OR timestamp = '' OR substr(timestamp, 1, 10) <= {date})",
                f"(valid_from IS NULL OR valid_from = '' OR substr(valid_from, 1, 10) <= {date})",
                f"(valid_to IS NULL OR valid_to = '' OR substr(valid_to, 1, 10) > {date})",
            )
        )
    else:
        clauses.append("(status IS NULL OR status = '' OR lower(status) = 'active')")
    if since:
        date = _sql_literal(since[:10])
        clauses.append(
            "(coalesce(nullif(timestamp, ''), valid_from) IS NULL OR "
            "coalesce(nullif(timestamp, ''), valid_from) = '' OR "
            f"substr(coalesce(nullif(timestamp, ''), valid_from), 1, 10) >= {date})"
        )
    return " AND ".join(clauses)


def upsert_vectors(
    paths: list[str],
    titles: list[str],
    summaries: list[str],
    projects: list[str],
    timestamps: list[str],
    vectors: list[list[float]],
    model: str = "BAAI/bge-small-en-v1.5",
) -> int:
    """Closed: destructive live-table mutation is no longer a public path.

    Publish only via publish_generation_vectors() into an immutable generation
    directory, then activate through the generation catalog CAS.
    """
    raise RuntimeError(
        "destructive Lance live-table upsert is closed; "
        "publish vectors into an immutable generation via "
        "publish_generation_vectors() and activate with catalog CAS"
    )


def publish_generation_vectors(
    *,
    generation_dir: Path,
    paths: list[str],
    titles: list[str],
    summaries: list[str],
    projects: list[str],
    timestamps: list[str],
    vectors: list[list[float]],
    model: str = "BAAI/bge-small-en-v1.5",
) -> dict[str, object]:
    """Write Lance vectors under an immutable generation directory only.

    Never drops or mutates the legacy live TABLE_NAME. Activation is owned by
    the generation catalog CAS outside this helper.
    """
    directory = Path(generation_dir)
    try:
        directory.mkdir(parents=True, exist_ok=True)
        if not _have_lancedb():
            return {"status": "skipped", "reason": "lancedb_unavailable", "count": 0}
        import lancedb
        import pyarrow as pa

        lance_dir = directory / "lance"
        if lance_dir.exists() or lance_dir.is_symlink():
            return {"status": "skipped", "reason": "already_exists", "count": 0}
        db = lancedb.connect(str(lance_dir))
        schema = pa.schema(
            [
                pa.field("path", pa.string()),
                pa.field("title", pa.string()),
                pa.field("summary", pa.string()),
                pa.field("project", pa.string()),
                pa.field("timestamp", pa.string()),
                pa.field("vector", pa.list_(pa.float32(), EMBEDDING_DIM)),
                pa.field("model", pa.string()),
            ]
        )
        data = pa.table(
            {
                "path": paths,
                "title": titles,
                "summary": summaries,
                "project": projects,
                "timestamp": timestamps,
                "vector": vectors,
                "model": [model] * len(vectors),
            },
            schema=schema,
        )
        table = db.create_table("chunks", data=data)
        if table.count_rows() != len(vectors):
            return {"status": "failed", "reason": "count_mismatch", "count": 0}
        return {"status": "ok", "reason": None, "count": len(vectors), "path": str(lance_dir)}
    except Exception as exc:
        return {"status": "failed", "reason": f"{type(exc).__name__}", "count": 0}


def vector_search(
    query_vec: list[float],
    limit: int = 20,
    project: str | None = None,
    since: str | None = None,
    as_of: str | None = None,
    *,
    scope: str = "all",
    authority: str | None = None,
    expected_model_id: str = EMBEDDING_MODEL,
    expected_model_revision: str = EMBEDDING_MODEL_REVISION,
    expected_sources: list[tuple[str, str]] | None = None,
) -> list[dict]:
    """Search vectors via LanceDB. Returns ranked larger-is-better similarities."""
    db = _get_db()
    if db is None:
        return []

    try:
        if expected_sources is None or (project and not _PROJECT_SAFE.fullmatch(project)):
            return []
        table = db.open_table(TABLE_NAME)
        membership = (
            table.query()
            .select(["path", "source_sha256", "model", "model_revision"])
            .to_list()
        )
        actual_sources = [
            (str(row.get("path") or ""), str(row.get("source_sha256") or ""))
            for row in membership
        ]
        if (
            len(actual_sources) != len(expected_sources)
            or len(dict(actual_sources)) != len(actual_sources)
            or dict(actual_sources) != dict(expected_sources)
            or any(row.get("model") != expected_model_id for row in membership)
            or any(
                row.get("model_revision") != expected_model_revision
                for row in membership
            )
        ):
            return []
        query = table.search(query_vec)
        query = query.where(
            _lance_filter(
                project=project,
                since=since,
                as_of=as_of,
                scope=scope,
                authority=authority,
                model_id=expected_model_id,
                model_revision=expected_model_revision,
            )
        )
        query = query.limit(limit)
        raw = query.to_list()
        rows = _rows_from_lance_hits(raw)
        rows = apply_vector_filters(
            rows,
            project=project,
            since=since,
            as_of=as_of,
            scope=scope,
            authority=authority,
        )
        return rows[:limit]
    except Exception:
        return []


def vector_count() -> int:
    """Return number of vectors in LanceDB, or 0 if unavailable."""
    if not have_lancedb():
        return 0
    try:
        db = _get_db()
        table = db.open_table(TABLE_NAME)
        return table.count_rows()
    except Exception:
        return 0


if __name__ == "__main__":
    if _have_lancedb():
        count = vector_count()
        print(f"LanceDB: {count} vectors in {LANCEDB_DIR}")
    else:
        print("LanceDB not installed. Run: uv sync --extra hybrid")
