"""LanceDB vector store for RAG."""

import logging
from pathlib import Path
from typing import Any

import lancedb
from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)

# Default embedding model - small but effective for code
DEFAULT_MODEL = "all-MiniLM-L6-v2"


class RAGStore:
    """Vector store using LanceDB for semantic code search."""

    def __init__(
        self,
        db_path: str | Path | None = None,
        model_name: str = DEFAULT_MODEL,
    ):
        """Initialize RAG store.

        Args:
            db_path: Path to LanceDB database. Defaults to ~/.manimgl-mcp/rag.db
            model_name: Sentence transformer model for embeddings
        """
        if db_path is None:
            db_path = Path.home() / ".manimgl-mcp" / "rag.db"

        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        self.model_name = model_name
        self._model: SentenceTransformer | None = None
        self._db: lancedb.DBConnection | None = None

        # Table names
        self.EXAMPLES_TABLE = "scene_examples"
        self.API_TABLE = "api_signatures"
        self.PATTERNS_TABLE = "animation_patterns"
        self.ERRORS_TABLE = "error_patterns"
        self.DOCS_TABLE = "documentation"

    @property
    def model(self) -> SentenceTransformer:
        """Lazy-load embedding model."""
        if self._model is None:
            logger.info(f"Loading embedding model: {self.model_name}")
            self._model = SentenceTransformer(self.model_name)
        return self._model

    @property
    def db(self) -> lancedb.DBConnection:
        """Lazy-load database connection."""
        if self._db is None:
            logger.info(f"Connecting to LanceDB: {self.db_path}")
            self._db = lancedb.connect(str(self.db_path))
        return self._db

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for texts."""
        return self.model.encode(texts, convert_to_numpy=True).tolist()

    def embed_single(self, text: str) -> list[float]:
        """Generate embedding for single text."""
        return self.embed([text])[0]

    def create_table(
        self,
        name: str,
        data: list[dict[str, Any]],
        mode: str = "overwrite",
    ) -> lancedb.table.Table:
        """Create or overwrite a table with data.

        Args:
            name: Table name
            data: List of dicts, each must have 'vector' and 'text' keys
            mode: 'overwrite' or 'append'
        """
        logger.info(f"Creating table '{name}' with {len(data)} records")
        return self.db.create_table(name, data, mode=mode)

    def add_to_table(self, name: str, data: list[dict[str, Any]]) -> None:
        """Add records to existing table."""
        table = self.db.open_table(name)
        table.add(data)
        logger.info(f"Added {len(data)} records to '{name}'")

    def search(
        self,
        table_name: str,
        query: str,
        limit: int = 10,
        filter_expr: str | None = None,
    ) -> list[dict[str, Any]]:
        """Search table with semantic similarity.

        Args:
            table_name: Name of table to search
            query: Search query
            limit: Max results
            filter_expr: Optional SQL-like filter

        Returns:
            List of matching records with _distance score
        """
        try:
            table = self.db.open_table(table_name)
            query_vector = self.embed_single(query)

            search = table.search(query_vector).limit(limit)
            if filter_expr:
                search = search.where(filter_expr)

            return search.to_list()
        except Exception as e:
            logger.error(f"Search failed: {e}")
            return []

    def hybrid_search(
        self,
        table_name: str,
        query: str,
        limit: int = 10,
        alpha: float = 0.5,  # noqa: ARG002 - Reserved for weighted hybrid search
    ) -> list[dict[str, Any]]:
        """Hybrid search combining vector + full-text (BM25).

        Args:
            table_name: Table to search
            query: Search query
            limit: Max results
            alpha: Balance between vector (0) and text (1) search
        """
        try:
            table = self.db.open_table(table_name)
            query_vector = self.embed_single(query)

            # LanceDB hybrid search
            results = table.search(query_vector, query_type="hybrid").limit(limit).to_list()
            return results
        except Exception as e:
            # Fall back to vector-only if hybrid not supported
            logger.warning(f"Hybrid search failed, using vector: {e}")
            return self.search(table_name, query, limit)

    def table_exists(self, name: str) -> bool:
        """Check if table exists."""
        return name in self.db.table_names()

    def get_stats(self) -> dict[str, Any]:
        """Get database statistics."""
        stats = {"tables": {}}
        for name in self.db.table_names():
            table = self.db.open_table(name)
            stats["tables"][name] = {
                "count": table.count_rows(),
            }
        return stats

    def close(self) -> None:
        """Close database connection."""
        self._db = None
        self._model = None


# Global store instance
_store: RAGStore | None = None


def get_store() -> RAGStore:
    """Get or create global RAG store instance."""
    global _store
    if _store is None:
        _store = RAGStore()
    return _store
