"""Vector search abstraction for GreyMatter.

Provides a Protocol-based interface with SQLite-vec (Solo mode) and
pgvector (Distributed mode) implementations. Both use cosine distance
and float32[768] embeddings.

Usage:
    # Solo mode
    vs = SqliteVecSearch(db)
    vs.upsert("vec_observations", "obs-123", embedding)
    results = vs.search("vec_observations", query_embedding, limit=10)

    # Distributed mode
    vs = PgVectorSearch(crdb_pool)
    vs.upsert("vec_observations", "obs-123", embedding)
    results = vs.search("vec_observations", query_embedding, limit=10)
"""
import struct
from typing import Any, Protocol, runtime_checkable


def serialize_f32(vector: list[float]) -> bytes:
    """Serialize float list to bytes for sqlite-vec."""
    return struct.pack(f"{len(vector)}f", *vector)


@runtime_checkable
class VectorSearch(Protocol):
    """Abstract vector search interface.

    Implementations provide cosine-distance similarity search over
    float32 embeddings. Tables must be pre-created (via schema module).
    """

    def upsert(self, table: str, id: str, embedding: list[float]) -> None:
        """Insert or update an embedding vector."""
        ...

    def search(
        self, table: str, query_embedding: list[float], limit: int = 10
    ) -> list[dict[str, Any]]:
        """Search for nearest neighbors by cosine distance.

        Returns list of dicts with 'id' and 'distance' keys, ordered
        by ascending distance (closest first).
        """
        ...

    def delete(self, table: str, id: str) -> None:
        """Delete an embedding by ID."""
        ...

    def count(self, table: str) -> int:
        """Return number of indexed vectors in a table."""
        ...


class SqliteVecSearch:
    """sqlite-vec implementation for Solo mode.

    Uses vec0 virtual tables with cosine distance. Requires sqlite-vec
    extension loaded on the connection (done by SQLiteDatabase.__init__).

    The vec0 virtual table schema is:
        CREATE VIRTUAL TABLE {table} USING vec0(
            {pk_col} TEXT PRIMARY KEY,
            embedding float[768]
        )
    """

    def __init__(self, conn: Any):
        """Initialize with a sqlite3 connection (vec extension loaded)."""
        self._conn = conn

    def upsert(self, table: str, id: str, embedding: list[float]) -> None:
        """Insert or replace a vector.

        vec0 virtual tables don't support INSERT OR REPLACE, so we
        delete-then-insert to achieve upsert semantics.
        """
        blob = serialize_f32(embedding)
        pk_col = self._get_pk_col(table)
        self._conn.execute(f"DELETE FROM {table} WHERE {pk_col} = ?", (id,))
        self._conn.execute(
            f"INSERT INTO {table} ({pk_col}, embedding) VALUES (?, ?)",
            (id, blob),
        )
        self._conn.commit()

    def search(
        self, table: str, query_embedding: list[float], limit: int = 10
    ) -> list[dict[str, Any]]:
        """Search using vec_distance_cosine."""
        blob = serialize_f32(query_embedding)
        pk_col = self._get_pk_col(table)
        rows = self._conn.execute(
            f"SELECT {pk_col}, distance FROM {table} "
            f"WHERE embedding MATCH ? AND k = ? "
            f"ORDER BY distance",
            (blob, limit),
        ).fetchall()
        return [{"id": row[0], "distance": row[1]} for row in rows]

    def delete(self, table: str, id: str) -> None:
        """Delete a vector by primary key."""
        pk_col = self._get_pk_col(table)
        self._conn.execute(f"DELETE FROM {table} WHERE {pk_col} = ?", (id,))
        self._conn.commit()

    def count(self, table: str) -> int:
        """Count vectors in a table."""
        row = self._conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()
        return row[0] if row else 0

    @staticmethod
    def _get_pk_col(table: str) -> str:
        """Derive primary key column name from table name.

        Convention: vec_observations → observation_id,
                    vec_documents → document_id, etc.
        """
        # Strip 'vec_' prefix and singularize
        name = table.removeprefix("vec_")
        if name.endswith("s"):
            name = name[:-1]
        return f"{name}_id"


class PgVectorSearch:
    """pgvector implementation for Distributed mode (CockroachDB v25.1+).

    Uses VECTOR(768) columns with <=> cosine distance operator.
    CockroachDB v25.1+ has native pgvector compatibility.

    The pgvector table schema is:
        CREATE TABLE {table} (
            id TEXT PRIMARY KEY,
            embedding VECTOR(768)
        )
    """

    def __init__(self, get_conn: Any):
        """Initialize with a connection getter (callable or pool).

        get_conn should be a callable that returns a psycopg2 connection.
        The caller is responsible for connection lifecycle.
        """
        self._get_conn = get_conn

    def upsert(self, table: str, id: str, embedding: list[float]) -> None:
        """Insert or update a vector using ON CONFLICT."""
        vec_str = "[" + ",".join(str(f) for f in embedding) + "]"
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    f"INSERT INTO {table} (id, embedding) VALUES (%s, %s::vector) "
                    f"ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding",
                    (id, vec_str),
                )
            conn.commit()
        except Exception:
            conn.rollback()
            raise

    def search(
        self, table: str, query_embedding: list[float], limit: int = 10
    ) -> list[dict[str, Any]]:
        """Search using <=> cosine distance operator."""
        vec_str = "[" + ",".join(str(f) for f in query_embedding) + "]"
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    f"SELECT id, embedding <=> %s::vector AS distance "
                    f"FROM {table} ORDER BY distance LIMIT %s",
                    (vec_str, limit),
                )
                rows = cur.fetchall()
                return [{"id": row[0], "distance": row[1]} for row in rows]
        finally:
            pass  # Connection returned by caller

    def delete(self, table: str, id: str) -> None:
        """Delete a vector by ID."""
        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute(f"DELETE FROM {table} WHERE id = %s", (id,))
            conn.commit()
        except Exception:
            conn.rollback()
            raise

    def count(self, table: str) -> int:
        """Count vectors in a table."""
        conn = self._get_conn()
        with conn.cursor() as cur:
            cur.execute(f"SELECT COUNT(*) FROM {table}")
            row = cur.fetchone()
            return row[0] if row else 0
