"""SQLite vector backend — stores embeddings as BLOBs, computes cosine similarity in Python.

Uses a regular SQLite table (not sqlite-vec's vec0 virtual table) so it works
on every platform regardless of extension-loading support. At Sulci's scale
(hundreds to low thousands of atoms) this is perfectly fast.
"""

import logging
import math
import struct

from sqlalchemy import text

logger = logging.getLogger(__name__)


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two vectors. Returns 0-1 range."""
    dot = sum(x * y for x, y in zip(a, b, strict=False))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


class SqliteVecBackend:
    """Vector storage using a regular SQLite table with BLOB embeddings.

    Shares the same SQLAlchemy session_factory (and therefore the same
    SQLite connection via StaticPool) as the relational data — no WAL
    conflicts, no dual-connection issues.
    """

    def __init__(self, session_factory, dimensions: int = 512) -> None:
        self._session_factory = session_factory
        self._dimensions = dimensions

    async def initialize(self) -> None:
        """Create the atom_vectors table if it doesn't exist."""
        async with self._session_factory() as session, session.begin():
            await session.execute(
                text("CREATE TABLE IF NOT EXISTS atom_vectors (  atom_id TEXT PRIMARY KEY,  embedding BLOB NOT NULL)")
            )
        logger.info("SQLite vector backend initialized (dimensions=%d)", self._dimensions)

    @staticmethod
    def _serialize_embedding(embedding: list[float]) -> bytes:
        """Pack a float list into a compact binary blob."""
        return struct.pack(f"{len(embedding)}f", *embedding)

    @staticmethod
    def _deserialize_embedding(blob: bytes) -> list[float]:
        """Unpack a binary blob into a float list."""
        count = len(blob) // 4
        return list(struct.unpack(f"{count}f", blob))

    async def upsert(self, atom_id: str, embedding: list[float], metadata: dict, document: str) -> None:
        """Insert or replace a single vector."""
        blob = self._serialize_embedding(embedding)
        async with self._session_factory() as session, session.begin():
            await session.execute(
                text(
                    "INSERT INTO atom_vectors (atom_id, embedding) VALUES (:id, :emb) "
                    "ON CONFLICT(atom_id) DO UPDATE SET embedding = :emb"
                ),
                {"id": atom_id, "emb": blob},
            )

    async def upsert_batch(self, items: list[tuple[str, list[float], dict, str]]) -> None:
        """Bulk insert vectors."""
        if not items:
            return
        async with self._session_factory() as session, session.begin():
            for atom_id, embedding, _metadata, _document in items:
                blob = self._serialize_embedding(embedding)
                await session.execute(
                    text(
                        "INSERT INTO atom_vectors (atom_id, embedding) VALUES (:id, :emb) "
                        "ON CONFLICT(atom_id) DO UPDATE SET embedding = :emb"
                    ),
                    {"id": atom_id, "emb": blob},
                )

    async def delete(self, atom_id: str) -> None:
        """Remove a vector by atom ID."""
        async with self._session_factory() as session, session.begin():
            await session.execute(text("DELETE FROM atom_vectors WHERE atom_id = :id"), {"id": atom_id})

    async def delete_batch(self, atom_ids: list[str]) -> None:
        """Remove multiple vectors by atom IDs."""
        if not atom_ids:
            return
        async with self._session_factory() as session, session.begin():
            for atom_id in atom_ids:
                await session.execute(text("DELETE FROM atom_vectors WHERE atom_id = :id"), {"id": atom_id})

    async def update_metadata(self, atom_id: str, metadata: dict) -> None:
        """No-op — metadata lives in the relational knowledge_atoms table."""
        pass

    async def search(
        self, query_embedding: list[float], limit: int = 10, where: dict | None = None
    ) -> list[tuple[str, float]]:
        """KNN cosine search.

        Fetches all vectors from the DB, computes cosine similarity in Python,
        and returns the top-K results. When a where filter is provided, JOINs
        with knowledge_atoms for filtering.

        Returns list of (atom_id, cosine_distance) where distance is in [0, 2].
        """
        async with self._session_factory() as session:
            if where:
                conditions = []
                params: dict = {}
                for key, value in where.items():
                    if key == "is_active":
                        conditions.append(f"ka.is_active = :f_{key}")
                        params[f"f_{key}"] = 1 if value else 0
                    elif key == "atom_type":
                        conditions.append(f"ka.atom_type = :f_{key}")
                        params[f"f_{key}"] = value
                where_clause = " AND ".join(conditions) if conditions else "1=1"

                sql = (
                    f"SELECT v.atom_id, v.embedding "
                    f"FROM atom_vectors v "
                    f"INNER JOIN knowledge_atoms ka ON ka.id = v.atom_id "
                    f"WHERE {where_clause}"
                )
            else:
                params = {}
                sql = "SELECT atom_id, embedding FROM atom_vectors"

            result = await session.execute(text(sql), params)
            rows = result.fetchall()

        # Compute cosine similarity in Python and sort
        scored: list[tuple[str, float]] = []
        for row in rows:
            atom_id = row[0]
            emb = self._deserialize_embedding(row[1])
            sim = _cosine_similarity(query_embedding, emb)
            # Convert similarity to distance: distance = 2 * (1 - similarity)
            distance = 2.0 * (1.0 - sim)
            scored.append((atom_id, distance))

        scored.sort(key=lambda x: x[1])
        return scored[:limit]

    async def get_all_with_embeddings(self, where: dict | None = None, limit: int = 500) -> dict:
        """Get all vectors with embeddings for duplicate detection."""
        async with self._session_factory() as session:
            if where:
                conditions = []
                params: dict = {"lim": limit}
                for key, value in where.items():
                    if key == "is_active":
                        conditions.append(f"ka.is_active = :f_{key}")
                        params[f"f_{key}"] = 1 if value else 0
                    elif key == "atom_type":
                        conditions.append(f"ka.atom_type = :f_{key}")
                        params[f"f_{key}"] = value
                where_clause = " AND ".join(conditions) if conditions else "1=1"

                sql = (
                    f"SELECT v.atom_id, v.embedding "
                    f"FROM atom_vectors v "
                    f"INNER JOIN knowledge_atoms ka ON ka.id = v.atom_id "
                    f"WHERE {where_clause} "
                    f"LIMIT :lim"
                )
            else:
                params = {"lim": limit}
                sql = "SELECT atom_id, embedding FROM atom_vectors LIMIT :lim"

            result = await session.execute(text(sql), params)
            rows = result.fetchall()

        ids = []
        embeddings = []
        for row in rows:
            ids.append(row[0])
            emb = self._deserialize_embedding(row[1])
            embeddings.append(emb)

        return {"ids": ids, "embeddings": embeddings}

    def distance_to_similarity(self, distance: float) -> float:
        """Convert cosine distance to similarity.

        Distance = 2 * (1 - similarity), so similarity = 1 - distance/2.
        Range: [0, 1] where 1 = identical.
        """
        return 1.0 - (distance / 2.0)

    async def purge(self) -> None:
        """Delete all vectors."""
        async with self._session_factory() as session, session.begin():
            await session.execute(text("DELETE FROM atom_vectors"))
