"""Embedding store -- FAISS index + SQLite metadata for kNN routing data.

Stores embedding vectors alongside routing outcome metadata so the kNN
router can learn from past decisions.  Uses FAISS ``IndexFlatIP`` (inner
product on L2-normalised vectors = cosine similarity) for fast lookup and
SQLite for durable metadata.

Requires ``faiss-cpu`` and ``numpy`` (``[smart]`` tier).
"""

from __future__ import annotations

import logging
import sqlite3
import time
from typing import TYPE_CHECKING

import faiss
import numpy as np
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Guarded imports
# ---------------------------------------------------------------------------

FAISS_AVAILABLE: bool = False

try:
    import faiss
    import numpy as np

    FAISS_AVAILABLE = True
except ImportError:
    faiss = None  # type: ignore[assignment]
    np = None  # type: ignore[assignment]

if TYPE_CHECKING:
    from pathlib import Path


# ============================================================================
# Pydantic models
# ============================================================================


class EmbeddingMetadata(BaseModel):
    """Metadata stored alongside each embedding vector."""

    request_id: str
    model_used: str
    backend_type: str
    routing_tier: str  # "rule", "knn", "modernbert", "qwen"
    user_feedback: float | None = None  # 0.0-1.0 quality rating
    latency_ms: float = 0.0
    cost: float = 0.0
    was_exploration: bool = False  # 5% deliberate exploration
    timestamp: float = Field(default_factory=time.time)


class SearchResult(BaseModel):
    """A single kNN search result with similarity score and metadata."""

    score: float  # cosine similarity (higher = more similar)
    metadata: EmbeddingMetadata


class StoreStats(BaseModel):
    """Aggregate statistics about the embedding store."""

    total_embeddings: int
    index_size_bytes: int
    dimensions: int
    distinct_models: int = 0
    distinct_backends: int = 0
    avg_latency_ms: float = 0.0
    feedback_count: int = 0


# ============================================================================
# SQL schema
# ============================================================================

_SCHEMA_SQL = """\
CREATE TABLE IF NOT EXISTS embeddings (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    request_id  TEXT    NOT NULL,
    model_used  TEXT    NOT NULL,
    backend_type TEXT   NOT NULL,
    routing_tier TEXT   NOT NULL,
    user_feedback REAL,
    latency_ms  REAL    NOT NULL DEFAULT 0.0,
    cost        REAL    NOT NULL DEFAULT 0.0,
    was_exploration INTEGER NOT NULL DEFAULT 0,
    timestamp   REAL    NOT NULL,
    UNIQUE(request_id)
);

CREATE INDEX IF NOT EXISTS idx_embeddings_model ON embeddings(model_used);
CREATE INDEX IF NOT EXISTS idx_embeddings_backend ON embeddings(backend_type);
CREATE INDEX IF NOT EXISTS idx_embeddings_timestamp ON embeddings(timestamp);
"""


# ============================================================================
# EmbeddingStore
# ============================================================================


class EmbeddingStore:
    """Persistent vector store backed by FAISS and SQLite.

    The FAISS index stores L2-normalised 384-d vectors.  ``IndexFlatIP``
    (inner product) gives exact cosine similarity on normalised vectors.
    SQLite stores the routing metadata keyed by the same integer ID.

    Thread safety: FAISS ``IndexFlatIP`` supports concurrent reads.  Writes
    (``add``) must be serialised by the caller (the kNN router does this
    via ``asyncio`` which is single-threaded per loop).
    """

    def __init__(self, index_path: Path, db_path: Path, dim: int = 384) -> None:
        if not FAISS_AVAILABLE:
            raise RuntimeError(
                "EmbeddingStore requires faiss-cpu and numpy. Install them with: pip install llmhosts[smart]"
            )
        self._index_path = index_path
        self._db_path = db_path
        self._dim = dim
        self._index: faiss.IndexFlatIP | None = None
        self._db: sqlite3.Connection | None = None
        self._dirty: bool = False  # True when index has unsaved changes

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self) -> None:
        """Load or create the FAISS index and SQLite metadata store."""
        # Ensure parent directories exist
        self._index_path.parent.mkdir(parents=True, exist_ok=True)
        self._db_path.parent.mkdir(parents=True, exist_ok=True)

        # FAISS index
        if self._index_path.is_file():
            self._index = faiss.read_index(str(self._index_path))
            logger.info(
                "Loaded FAISS index from %s (%d vectors, %d dims)",
                self._index_path,
                self._index.ntotal,
                self._dim,
            )
        else:
            self._index = faiss.IndexFlatIP(self._dim)
            logger.info("Created new FAISS IndexFlatIP (dim=%d)", self._dim)

        # SQLite
        self._db = sqlite3.connect(str(self._db_path), check_same_thread=False)
        self._db.row_factory = sqlite3.Row
        self._db.executescript(_SCHEMA_SQL)
        self._db.commit()
        logger.info("Opened SQLite metadata store at %s", self._db_path)

        # Consistency check: FAISS ntotal should match SQLite row count
        db_count = self._db.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
        if self._index.ntotal != db_count:
            logger.warning(
                "FAISS/SQLite mismatch: index has %d vectors, DB has %d rows. Consider rebuilding the index.",
                self._index.ntotal,
                db_count,
            )

    # ------------------------------------------------------------------
    # Write operations
    # ------------------------------------------------------------------

    async def add(self, embedding: np.ndarray, metadata: EmbeddingMetadata) -> int:
        """Add an embedding with routing metadata.

        Parameters
        ----------
        embedding:
            L2-normalised vector, shape ``(dim,)`` or ``(1, dim)``.
        metadata:
            Routing outcome metadata for this request.

        Returns
        -------
        int
            The integer ID assigned to this embedding (1-indexed from SQLite).
        """
        if self._index is None or self._db is None:
            raise RuntimeError("Store not initialised -- call await store.initialize() first")

        # Ensure correct shape
        vec = embedding.reshape(1, -1).astype(np.float32)
        if vec.shape[1] != self._dim:
            raise ValueError(f"Expected {self._dim}-d vector, got {vec.shape[1]}-d")

        # Insert metadata into SQLite first to get the row ID
        cursor = self._db.execute(
            """
            INSERT OR REPLACE INTO embeddings
                (request_id, model_used, backend_type, routing_tier,
                 user_feedback, latency_ms, cost, was_exploration, timestamp)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                metadata.request_id,
                metadata.model_used,
                metadata.backend_type,
                metadata.routing_tier,
                metadata.user_feedback,
                metadata.latency_ms,
                metadata.cost,
                int(metadata.was_exploration),
                metadata.timestamp,
            ),
        )
        self._db.commit()
        row_id = cursor.lastrowid or 0

        # Add to FAISS index
        self._index.add(vec)
        self._dirty = True

        logger.debug(
            "Added embedding #%d for request %s (model=%s, backend=%s)",
            row_id,
            metadata.request_id,
            metadata.model_used,
            metadata.backend_type,
        )
        return row_id

    async def update_feedback(self, request_id: str, feedback: float) -> bool:
        """Update the user feedback score for an existing embedding.

        Returns ``True`` if the row was found and updated.
        """
        if self._db is None:
            raise RuntimeError("Store not initialised")

        cursor = self._db.execute(
            "UPDATE embeddings SET user_feedback = ? WHERE request_id = ?",
            (feedback, request_id),
        )
        self._db.commit()
        return cursor.rowcount > 0

    # ------------------------------------------------------------------
    # Search
    # ------------------------------------------------------------------

    async def search(self, query: np.ndarray, k: int = 10) -> list[SearchResult]:
        """Find the *k* nearest neighbours by cosine similarity.

        Parameters
        ----------
        query:
            L2-normalised query vector, shape ``(dim,)`` or ``(1, dim)``.
        k:
            Number of neighbours to return.

        Returns
        -------
        list[SearchResult]
            Sorted by descending similarity score.
        """
        if self._index is None or self._db is None:
            raise RuntimeError("Store not initialised")

        if self._index.ntotal == 0:
            return []

        vec = query.reshape(1, -1).astype(np.float32)
        actual_k = min(k, self._index.ntotal)

        # FAISS search returns (distances, indices) each shaped (1, k)
        scores, indices = self._index.search(vec, actual_k)

        results: list[SearchResult] = []
        for score, idx in zip(scores[0], indices[0], strict=False):
            if idx < 0:
                # FAISS returns -1 for unfilled slots
                continue

            # SQLite IDs are 1-indexed; FAISS IDs are 0-indexed
            row = self._db.execute("SELECT * FROM embeddings WHERE id = ?", (int(idx) + 1,)).fetchone()

            if row is None:
                logger.debug("No metadata for FAISS index %d -- skipping", idx)
                continue

            meta = EmbeddingMetadata(
                request_id=row["request_id"],
                model_used=row["model_used"],
                backend_type=row["backend_type"],
                routing_tier=row["routing_tier"],
                user_feedback=row["user_feedback"],
                latency_ms=row["latency_ms"],
                cost=row["cost"],
                was_exploration=bool(row["was_exploration"]),
                timestamp=row["timestamp"],
            )
            results.append(SearchResult(score=float(score), metadata=meta))

        return results

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    async def save(self) -> None:
        """Persist the FAISS index to disk."""
        if self._index is None:
            return
        if not self._dirty:
            logger.debug("FAISS index is clean -- skipping save")
            return

        faiss.write_index(self._index, str(self._index_path))
        self._dirty = False
        logger.info("Saved FAISS index to %s (%d vectors)", self._index_path, self._index.ntotal)

    async def close(self) -> None:
        """Save the index and close the SQLite connection."""
        await self.save()
        if self._db is not None:
            self._db.close()
            self._db = None
            logger.info("Closed embedding store")

    # ------------------------------------------------------------------
    # Statistics
    # ------------------------------------------------------------------

    async def get_stats(self) -> StoreStats:
        """Compute aggregate statistics about the store."""
        if self._index is None or self._db is None:
            return StoreStats(total_embeddings=0, index_size_bytes=0, dimensions=self._dim)

        total = self._index.ntotal

        # Estimate index size: IndexFlatIP stores float32 vectors
        index_bytes = total * self._dim * 4  # 4 bytes per float32

        # Aggregate from SQLite
        row = self._db.execute(
            """
            SELECT
                COUNT(DISTINCT model_used) AS distinct_models,
                COUNT(DISTINCT backend_type) AS distinct_backends,
                AVG(latency_ms) AS avg_latency,
                COUNT(user_feedback) AS feedback_count
            FROM embeddings
            """
        ).fetchone()

        return StoreStats(
            total_embeddings=total,
            index_size_bytes=index_bytes,
            dimensions=self._dim,
            distinct_models=row["distinct_models"] if row else 0,
            distinct_backends=row["distinct_backends"] if row else 0,
            avg_latency_ms=row["avg_latency"] or 0.0 if row else 0.0,
            feedback_count=row["feedback_count"] if row else 0,
        )

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def size(self) -> int:
        """Number of embeddings currently in the FAISS index."""
        if self._index is None:
            return 0
        return int(self._index.ntotal)

    @property
    def is_ready(self) -> bool:
        """``True`` when both FAISS and SQLite are initialised."""
        return self._index is not None and self._db is not None
