"""Vector store with FTS5 full-text search and embedding cache."""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import sqlite3
import struct
import threading
import time
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

# Default max cached embeddings before LRU eviction
DEFAULT_CACHE_MAX_ENTRIES = 10_000


@dataclass
class SearchResult:
    """A single search result from the vector store."""

    chunk_id: str
    text: str
    metadata: dict
    score: float  # 0..1, higher is better


class VectorStore(ABC):
    """Abstract base for vector stores."""

    @abstractmethod
    def add(
        self,
        texts: list[str],
        embeddings: list[list[float]],
        *,
        source: str = "",
        user_id: str = "",
        metadatas: list[dict] | None = None,
    ) -> list[str]:
        """Add chunks with embeddings. Returns list of chunk IDs."""
        ...

    @abstractmethod
    def search(
        self,
        query_embedding: list[float],
        *,
        top_k: int = 5,
        user_id: str | None = None,
        source: str | None = None,
    ) -> list[SearchResult]:
        """Search for similar chunks."""
        ...

    @abstractmethod
    def search_fts(self, query: str, *, top_k: int = 5) -> list[SearchResult]:
        """Full-text search using FTS5."""
        ...

    @abstractmethod
    def delete_source(self, source: str) -> int:
        """Delete all chunks from a source. Returns count deleted."""
        ...

    @abstractmethod
    def list_sources(self) -> list[dict]:
        """List distinct sources with chunk counts."""
        ...

    async def async_add(
        self,
        texts: list[str],
        embeddings: list[list[float]],
        **kwargs,
    ) -> list[str]:
        """Async wrapper for add()."""
        return await asyncio.to_thread(self.add, texts, embeddings, **kwargs)

    async def async_search(
        self,
        query_embedding: list[float],
        **kwargs,
    ) -> list[SearchResult]:
        """Async wrapper for search()."""
        return await asyncio.to_thread(self.search, query_embedding, **kwargs)

    async def async_search_fts(self, query: str, **kwargs) -> list[SearchResult]:
        """Async wrapper for search_fts()."""
        return await asyncio.to_thread(self.search_fts, query, **kwargs)


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


def _hash_text(text: str) -> str:
    """Hash text content for embedding cache key."""
    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]


def _build_fts_query(raw: str) -> str | None:
    """Build FTS5 query from raw text. Returns None if no valid tokens."""
    import re
    tokens = re.findall(r"[\w]+", raw, re.UNICODE)
    if not tokens:
        return None
    # Quote each token and AND them together
    quoted = [f'"{t}"' for t in tokens if len(t) > 1]
    if not quoted:
        return None
    return " AND ".join(quoted)


def _bm25_rank_to_score(rank: float) -> float:
    """Convert SQLite FTS5 bm25() rank to 0..1 score.

    FTS5 bm25() returns negative values where more negative = more relevant.
    """
    if not isinstance(rank, (int, float)):
        return 0.01
    if rank < 0:
        relevance = -rank
        return relevance / (1.0 + relevance)
    return 1.0 / (1.0 + rank)


class SqliteVecStore(VectorStore):
    """SQLite-backed vector store with FTS5 and embedding cache."""

    def __init__(
        self,
        db_path: str,
        dimensions: int = 768,
        cache_max_entries: int = DEFAULT_CACHE_MAX_ENTRIES,
    ):
        self.db_path = db_path
        self.dimensions = dimensions
        self.cache_max_entries = cache_max_entries
        self._conn: sqlite3.Connection | None = None
        self._vec_available = False
        self._fts_available = False
        self._lock = threading.Lock()
        self._init_db()

    def _init_db(self) -> None:
        """Initialize the database schema: chunks, vec, FTS5, embedding cache."""
        self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")

        # Load sqlite-vec extension
        try:
            import sqlite_vec

            self._conn.enable_load_extension(True)
            sqlite_vec.load(self._conn)
            self._conn.enable_load_extension(False)
            self._vec_available = True
            logger.info("sqlite-vec extension loaded successfully")
        except (ImportError, OSError) as exc:
            logger.warning(
                "sqlite-vec not available (%s). "
                "Vector search will be disabled. "
                "Install with: pip install sqlite-vec",
                exc,
            )

        # Create chunks metadata table
        self._conn.execute("""
            CREATE TABLE IF NOT EXISTS chunks (
                id TEXT PRIMARY KEY,
                text TEXT NOT NULL,
                source TEXT NOT NULL DEFAULT '',
                user_id TEXT NOT NULL DEFAULT '',
                metadata TEXT NOT NULL DEFAULT '{}',
                created_at REAL NOT NULL
            )
        """)

        # Create virtual table for vector search
        if self._vec_available:
            self._conn.execute(f"""
                CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec
                USING vec0(embedding float[{self.dimensions}])
            """)

        # FTS5 full-text search index
        try:
            self._conn.execute("""
                CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
                    text,
                    id UNINDEXED,
                    source UNINDEXED
                )
            """)
            self._fts_available = True
            logger.info("FTS5 full-text search index available")
        except sqlite3.OperationalError as exc:
            logger.warning("FTS5 not available (%s). Keyword search will use in-memory BM25.", exc)

        # Embedding cache table
        self._conn.execute("""
            CREATE TABLE IF NOT EXISTS embedding_cache (
                provider TEXT NOT NULL,
                model TEXT NOT NULL,
                hash TEXT NOT NULL,
                embedding BLOB NOT NULL,
                dims INTEGER NOT NULL,
                updated_at REAL NOT NULL,
                PRIMARY KEY (provider, model, hash)
            )
        """)
        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_cache_updated ON embedding_cache(updated_at)"
        )

        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_chunks_source ON chunks(source)"
        )
        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_chunks_user_id ON chunks(user_id)"
        )
        self._conn.commit()

    # ── Embedding cache ──

    def cache_get(
        self, hashes: list[str], provider: str, model: str
    ) -> dict[str, list[float]]:
        """Batch lookup embedding cache. Returns {hash: embedding} for hits."""
        if not hashes:
            return {}

        conn = self._conn
        assert conn is not None

        result: dict[str, list[float]] = {}
        # Query in batches of 400 (SQLite variable limit)
        for i in range(0, len(hashes), 400):
            batch = hashes[i : i + 400]
            placeholders = ",".join("?" for _ in batch)
            rows = conn.execute(
                f"SELECT hash, embedding, dims FROM embedding_cache "
                f"WHERE provider = ? AND model = ? AND hash IN ({placeholders})",
                [provider, model, *batch],
            ).fetchall()
            for h, blob, dims in rows:
                result[h] = list(struct.unpack(f"<{dims}f", blob))

        if result:
            logger.debug("Embedding cache: %d hits / %d queries", len(result), len(hashes))
        return result

    def cache_put(
        self,
        items: list[tuple[str, list[float]]],
        provider: str,
        model: str,
    ) -> None:
        """Batch insert/update embedding cache. items = [(hash, embedding), ...]."""
        if not items:
            return

        conn = self._conn
        assert conn is not None

        now = time.time()
        with self._lock:
            for h, embedding in items:
                dims = len(embedding)
                blob = struct.pack(f"<{dims}f", *embedding)
                conn.execute(
                    "INSERT INTO embedding_cache (provider, model, hash, embedding, dims, updated_at) "
                    "VALUES (?, ?, ?, ?, ?, ?) "
                    "ON CONFLICT(provider, model, hash) DO UPDATE SET "
                    "embedding=excluded.embedding, dims=excluded.dims, updated_at=excluded.updated_at",
                    (provider, model, h, blob, dims, now),
                )
            conn.commit()

            # LRU eviction (called under lock; _prune_cache must not re-acquire)
            self._prune_cache()

    def _prune_cache(self) -> None:
        """Evict oldest cache entries if over max_entries."""
        if self.cache_max_entries <= 0:
            return

        conn = self._conn
        assert conn is not None

        count = conn.execute("SELECT COUNT(*) FROM embedding_cache").fetchone()[0]
        excess = count - self.cache_max_entries
        if excess > 0:
            conn.execute(
                "DELETE FROM embedding_cache WHERE rowid IN ("
                "  SELECT rowid FROM embedding_cache ORDER BY updated_at ASC LIMIT ?"
                ")",
                (excess,),
            )
            conn.commit()
            logger.debug("Cache pruned: evicted %d oldest entries", excess)

    # ── Chunk operations ──

    def add(
        self,
        texts: list[str],
        embeddings: list[list[float]],
        *,
        source: str = "",
        user_id: str = "",
        metadatas: list[dict] | None = None,
    ) -> list[str]:
        """Add text chunks with their embeddings to the store."""
        if len(texts) != len(embeddings):
            raise ValueError(
                f"texts ({len(texts)}) and embeddings ({len(embeddings)}) must have same length"
            )

        if metadatas and len(metadatas) != len(texts):
            raise ValueError(
                f"metadatas ({len(metadatas)}) must match texts ({len(texts)})"
            )

        for i, emb in enumerate(embeddings):
            if len(emb) != self.dimensions:
                raise ValueError(
                    f"Embedding at index {i} has {len(emb)} dimensions, "
                    f"expected {self.dimensions}"
                )

        conn = self._conn
        assert conn is not None, "Database not initialized"

        now = time.time()
        chunk_ids: list[str] = []

        with self._lock:
            for i, (text, embedding) in enumerate(zip(texts, embeddings)):
                chunk_id = uuid.uuid4().hex[:16]
                meta = metadatas[i] if metadatas else {}
                meta_json = json.dumps(meta, ensure_ascii=False)

                cursor = conn.execute(
                    "INSERT INTO chunks (id, text, source, user_id, metadata, created_at) "
                    "VALUES (?, ?, ?, ?, ?, ?)",
                    (chunk_id, text, source, user_id, meta_json, now),
                )

                if self._vec_available:
                    conn.execute(
                        "INSERT INTO chunks_vec (rowid, embedding) VALUES (?, ?)",
                        (cursor.lastrowid, _serialize_f32(embedding)),
                    )

                # FTS5 index
                if self._fts_available:
                    conn.execute(
                        "INSERT INTO chunks_fts (text, id, source) VALUES (?, ?, ?)",
                        (text, chunk_id, source),
                    )

                chunk_ids.append(chunk_id)

            conn.commit()
        logger.debug("Added %d chunks from source=%r", len(chunk_ids), source)
        return chunk_ids

    def search(
        self,
        query_embedding: list[float],
        *,
        top_k: int = 5,
        user_id: str | None = None,
        source: str | None = None,
    ) -> list[SearchResult]:
        """Search for chunks similar to the query embedding."""
        if not self._vec_available:
            logger.warning("Vector search unavailable — sqlite-vec not loaded")
            return []

        conn = self._conn
        assert conn is not None, "Database not initialized"

        fetch_limit = top_k * 4 if (user_id or source) else top_k

        with self._lock:
            rows = conn.execute(
                "SELECT rowid, distance FROM chunks_vec "
                "WHERE embedding MATCH ? ORDER BY distance LIMIT ?",
                (_serialize_f32(query_embedding), fetch_limit),
            ).fetchall()

            if not rows:
                return []

            rowid_distance = {row[0]: row[1] for row in rows}
            rowid_list = list(rowid_distance.keys())

            placeholders = ",".join("?" for _ in rowid_list)
            chunk_rows = conn.execute(
                f"SELECT rowid, id, text, source, user_id, metadata, created_at FROM chunks "
                f"WHERE rowid IN ({placeholders})",
                rowid_list,
            ).fetchall()

        results: list[SearchResult] = []
        for crow in chunk_rows:
            c_rowid, c_id, c_text, c_source, c_user_id, meta_json, c_created_at = crow
            c_metadata = json.loads(meta_json)

            if user_id and c_user_id != user_id:
                continue
            if source and c_source != source:
                continue

            distance = rowid_distance.get(c_rowid, 1.0)
            score = 1.0 / (1.0 + distance)

            results.append(
                SearchResult(
                    chunk_id=c_id,
                    text=c_text,
                    metadata={
                        **c_metadata,
                        "source": c_source,
                        "user_id": c_user_id,
                        "created_at": c_created_at,
                    },
                    score=score,
                )
            )

        results.sort(key=lambda r: r.score, reverse=True)
        return results[:top_k]

    def search_fts(
        self,
        query: str,
        *,
        top_k: int = 5,
    ) -> list[SearchResult]:
        """Full-text search using FTS5 with BM25 ranking."""
        if not self._fts_available:
            return []

        conn = self._conn
        assert conn is not None

        fts_query = _build_fts_query(query)
        if not fts_query:
            return []

        try:
            with self._lock:
                rows = conn.execute(
                    "SELECT f.id, f.text, f.source, bm25(chunks_fts) as rank, c.created_at "
                    "FROM chunks_fts f LEFT JOIN chunks c ON f.id = c.id "
                    "WHERE chunks_fts MATCH ? "
                    "ORDER BY rank LIMIT ?",
                    (fts_query, top_k),
                ).fetchall()
        except sqlite3.OperationalError as e:
            logger.debug("FTS5 query failed: %s", e)
            return []

        return [
            SearchResult(
                chunk_id=chunk_id,
                text=text,
                metadata={"source": source, "created_at": created_at},
                score=_bm25_rank_to_score(rank),
            )
            for chunk_id, text, source, rank, created_at in rows
        ]

    @property
    def fts_available(self) -> bool:
        return self._fts_available

    def delete_source(self, source: str) -> int:
        """Delete all chunks from a given source."""
        conn = self._conn
        assert conn is not None, "Database not initialized"

        with self._lock:
            if self._vec_available:
                conn.execute(
                    "DELETE FROM chunks_vec WHERE rowid IN "
                    "(SELECT rowid FROM chunks WHERE source = ?)",
                    (source,),
                )
            if self._fts_available:
                conn.execute("DELETE FROM chunks_fts WHERE source = ?", (source,))

            cursor = conn.execute("DELETE FROM chunks WHERE source = ?", (source,))
            count = cursor.rowcount
            conn.commit()
        logger.debug("Deleted %d chunks from source=%r", count, source)
        return count

    def list_sources(self) -> list[dict]:
        """List distinct sources with chunk counts."""
        conn = self._conn
        assert conn is not None, "Database not initialized"

        with self._lock:
            rows = conn.execute(
                "SELECT source, COUNT(*) as chunk_count, MIN(created_at) as first_added "
                "FROM chunks GROUP BY source ORDER BY source"
            ).fetchall()

        return [
            {
                "source": row[0],
                "chunk_count": row[1],
                "first_added": row[2],
            }
            for row in rows
        ]

    def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None
