"""
Vector Store Interface

Provides storage and retrieval for document chunks using Qdrant.
Supports hybrid search combining dense (semantic) and sparse (keyword) vectors.
"""

import logging
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class SearchResult:
    """Result from vector search."""

    chunk_id: str
    document_id: str
    doc_id: str  # Human-readable ID (e.g., "ANSI/ANS-8.1-2014")
    content: str
    score: float
    section_path: str = ""
    page_number: int | None = None
    chunk_index: int = 0
    has_table: bool = False
    has_equation: bool = False
    source: str = ""
    metadata: dict = field(default_factory=dict)

    @property
    def citation(self) -> str:
        """Generate citation string for this result."""
        parts = [self.doc_id]
        if self.section_path:
            parts.append(self.section_path)
        if self.page_number:
            parts.append(f"p. {self.page_number}")
        return ", ".join(parts)


class QdrantMode(Enum):
    """Qdrant deployment mode."""
    LOCAL = "local"
    CLOUD = "cloud"
    MEMORY = "memory"  # In-memory for testing


@dataclass
class VectorStoreConfig:
    """Configuration for vector store."""

    # Qdrant settings
    collection_name: str = "sage_ncs_chunks"
    mode: QdrantMode = QdrantMode.LOCAL

    # Local mode settings
    host: str = "localhost"
    port: int = 6333
    grpc_port: int = 6334
    prefer_grpc: bool = True

    # Cloud mode settings
    cloud_url: str | None = None
    api_key: str | None = None

    # Vector settings
    dense_vector_size: int = 1536  # OpenAI text-embedding-3-small
    distance_metric: str = "Cosine"

    # Sparse vector settings (for hybrid search)
    enable_sparse: bool = True

    # Index settings
    on_disk: bool = False
    hnsw_m: int = 16
    hnsw_ef_construct: int = 100

    # Payload indexing for efficient filtering
    indexed_payload_fields: list[str] = field(default_factory=lambda: [
        "doc_id",
        "document_id",
        "source",
        "chunk_type",
        "has_table",
        "has_equation",
    ])


class QdrantVectorStore:
    """
    Vector database interface for SAGE using Qdrant.

    Stores document chunks with their embeddings and metadata,
    enabling efficient semantic search. Supports both local
    Qdrant instances and Qdrant Cloud.

    Collection Schema:
        - Dense vectors: 1536 dimensions (text-embedding-3-small)
        - Sparse vectors: For hybrid BM25-style search
        - Payload fields:
            - chunk_id: Unique chunk identifier
            - doc_id: Human-readable document ID (e.g., "ANSI/ANS-8.1-2014")
            - document_id: Internal document UUID
            - content: Full chunk text
            - section_path: Hierarchical section path
            - page_number: Source page number
            - chunk_index: Position in document
            - has_table: Whether chunk contains a table
            - has_equation: Whether chunk contains an equation
            - source: Document source (e.g., "ansi", "nureg", "osti")
            - chunk_type: Type of content (text, table, equation, etc.)

    Example:
        # Local mode
        config = VectorStoreConfig(mode=QdrantMode.LOCAL)
        store = QdrantVectorStore(config)
        await store.initialize()

        # Cloud mode
        config = VectorStoreConfig(
            mode=QdrantMode.CLOUD,
            cloud_url="https://your-cluster.qdrant.io",
            api_key="your-api-key",
        )
        store = QdrantVectorStore(config)
        await store.initialize()

        # Add chunks
        await store.add_chunks(chunks, embeddings)

        # Search
        results = await store.search(query_embedding, top_k=10)

        # Hybrid search
        results = await store.hybrid_search(
            dense_embedding=dense_vec,
            sparse_embedding=sparse_vec,
            top_k=10,
            filters={"source": "ansi"},
        )
    """

    def __init__(self, config: VectorStoreConfig | None = None):
        """
        Initialize vector store.

        Args:
            config: Vector store configuration
        """
        self.config = config or VectorStoreConfig()
        self._client = None
        self._async_client = None
        self._initialized = False

    async def initialize(self) -> None:
        """
        Initialize connection and create collection if needed.

        Creates the Qdrant client based on configuration mode,
        sets up the collection with proper vector configuration,
        and creates payload indices for efficient filtering.
        """
        try:
            from qdrant_client import AsyncQdrantClient, QdrantClient
            from qdrant_client.models import (  # noqa: F401
                Distance,
                HnswConfigDiff,
                OptimizersConfigDiff,
                PayloadSchemaType,
                VectorParams,
            )
        except ImportError:
            raise ImportError(
                "qdrant-client not installed. Run: pip install qdrant-client"
            ) from None

        # Create client based on mode
        if self.config.mode == QdrantMode.MEMORY:
            # In-memory mode for testing
            self._client = QdrantClient(":memory:")
            self._async_client = None  # Memory mode doesn't support async
            logger.info("Using in-memory Qdrant client")

        elif self.config.mode == QdrantMode.CLOUD:
            # Qdrant Cloud mode
            if not self.config.cloud_url:
                raise ValueError("cloud_url required for cloud mode")
            if not self.config.api_key:
                raise ValueError("api_key required for cloud mode")

            self._client = QdrantClient(
                url=self.config.cloud_url,
                api_key=self.config.api_key,
                prefer_grpc=self.config.prefer_grpc,
            )
            self._async_client = AsyncQdrantClient(
                url=self.config.cloud_url,
                api_key=self.config.api_key,
                prefer_grpc=self.config.prefer_grpc,
            )
            logger.info(f"Connected to Qdrant Cloud at {self.config.cloud_url}")

        else:
            # Local mode (default)
            self._client = QdrantClient(
                host=self.config.host,
                port=self.config.port,
                grpc_port=self.config.grpc_port,
                prefer_grpc=self.config.prefer_grpc,
            )
            self._async_client = AsyncQdrantClient(
                host=self.config.host,
                port=self.config.port,
                grpc_port=self.config.grpc_port,
                prefer_grpc=self.config.prefer_grpc,
            )
            logger.info(f"Connected to local Qdrant at {self.config.host}:{self.config.port}")

        # Check if collection exists
        collections = self._client.get_collections()
        collection_names = [c.name for c in collections.collections]

        if self.config.collection_name not in collection_names:
            await self._create_collection()
        else:
            logger.info(f"Using existing collection: {self.config.collection_name}")

        self._initialized = True
        logger.info("Vector store initialized")

    async def _create_collection(self) -> None:
        """Create collection with proper schema."""
        from qdrant_client.models import (
            Distance,
            HnswConfigDiff,
            OptimizersConfigDiff,
            SparseVectorParams,
            VectorParams,
        )

        distance = getattr(Distance, self.config.distance_metric.upper())

        # Configure dense vectors
        vectors_config = {
            "dense": VectorParams(
                size=self.config.dense_vector_size,
                distance=distance,
                on_disk=self.config.on_disk,
            ),
        }

        # Configure sparse vectors for hybrid search
        sparse_vectors_config = None
        if self.config.enable_sparse:
            sparse_vectors_config = {
                "sparse": SparseVectorParams(
                    index={"on_disk": self.config.on_disk},
                ),
            }

        # Create collection
        self._client.create_collection(
            collection_name=self.config.collection_name,
            vectors_config=vectors_config,
            sparse_vectors_config=sparse_vectors_config,
            optimizers_config=OptimizersConfigDiff(
                indexing_threshold=20000,
            ),
            hnsw_config=HnswConfigDiff(
                m=self.config.hnsw_m,
                ef_construct=self.config.hnsw_ef_construct,
            ),
        )
        logger.info(f"Created collection: {self.config.collection_name}")

        # Create payload indices for efficient filtering
        await self._create_payload_indices()

    async def _create_payload_indices(self) -> None:
        """Create indices on payload fields for efficient filtering."""
        from qdrant_client.models import PayloadSchemaType, TextIndexParams, TokenizerType

        # Define field types for indexing
        field_types = {
            "doc_id": PayloadSchemaType.KEYWORD,
            "document_id": PayloadSchemaType.KEYWORD,
            "source": PayloadSchemaType.KEYWORD,
            "chunk_type": PayloadSchemaType.KEYWORD,
            "has_table": PayloadSchemaType.BOOL,
            "has_equation": PayloadSchemaType.BOOL,
            "page_number": PayloadSchemaType.INTEGER,
            "chunk_index": PayloadSchemaType.INTEGER,
        }

        for field_name in self.config.indexed_payload_fields:
            if field_name in field_types:
                try:
                    self._client.create_payload_index(
                        collection_name=self.config.collection_name,
                        field_name=field_name,
                        field_schema=field_types[field_name],
                    )
                    logger.debug(f"Created payload index for: {field_name}")
                except Exception as e:
                    logger.warning(f"Could not create index for {field_name}: {e}")

        # Create full-text index on content for keyword search fallback
        try:
            self._client.create_payload_index(
                collection_name=self.config.collection_name,
                field_name="content",
                field_schema=TextIndexParams(
                    type="text",
                    tokenizer=TokenizerType.WORD,
                    min_token_len=2,
                    max_token_len=20,
                    lowercase=True,
                ),
            )
            logger.debug("Created full-text index on content")
        except Exception as e:
            logger.warning(f"Could not create full-text index: {e}")

    @property
    def client(self):
        """Get synchronous Qdrant client."""
        if not self._initialized:
            raise RuntimeError("Vector store not initialized. Call initialize() first.")
        return self._client

    @property
    def async_client(self):
        """Get async Qdrant client (not available in memory mode)."""
        if not self._initialized:
            raise RuntimeError("Vector store not initialized. Call initialize() first.")
        if self._async_client is None:
            raise RuntimeError("Async client not available in memory mode")
        return self._async_client

    async def add_chunks(
        self,
        chunks: list,
        dense_embeddings: list[list[float]],
        sparse_embeddings: list[dict] | None = None,
        batch_size: int = 100,
    ) -> list[str]:
        """
        Add chunks with their embeddings to the store.

        Args:
            chunks: List of Chunk objects (from chunking module)
            dense_embeddings: Dense embedding vectors (1536 dims each)
            sparse_embeddings: Optional sparse embedding vectors for hybrid search
                Each dict should have 'indices' and 'values' keys
            batch_size: Number of points to upsert per batch

        Returns:
            List of point IDs for the added chunks

        Raises:
            ValueError: If number of chunks doesn't match embeddings
        """
        from qdrant_client.models import PointStruct, SparseVector

        if len(chunks) != len(dense_embeddings):
            raise ValueError(
                f"Number of chunks ({len(chunks)}) must match "
                f"number of embeddings ({len(dense_embeddings)})"
            )

        points = []
        point_ids = []

        for i, (chunk, dense_vec) in enumerate(zip(chunks, dense_embeddings, strict=True)):
            # Generate or use existing ID
            point_id = str(getattr(chunk, 'id', None) or uuid.uuid4())
            point_ids.append(point_id)

            # Build vectors dict
            vectors = {"dense": dense_vec}

            # Add sparse vector if provided
            if sparse_embeddings and i < len(sparse_embeddings):
                sparse = sparse_embeddings[i]
                if sparse.get("indices") and sparse.get("values"):
                    vectors["sparse"] = SparseVector(
                        indices=sparse["indices"],
                        values=sparse["values"],
                    )

            # Extract chunk type value
            chunk_type = getattr(chunk, 'chunk_type', 'text')
            if hasattr(chunk_type, 'value'):
                chunk_type = chunk_type.value

            # Build payload with all NCS-relevant fields
            payload = {
                "chunk_id": point_id,
                "document_id": str(getattr(chunk, 'document_id', '') or ''),
                "doc_id": str(getattr(chunk, 'doc_id', '') or ''),
                "content": str(getattr(chunk, 'content', '') or ''),
                "section_path": str(getattr(chunk, 'section_path', '') or ''),
                "chunk_index": int(getattr(chunk, 'chunk_index', 0) or 0),
                "page_number": getattr(chunk, 'page_number', None),
                "has_table": bool(getattr(chunk, 'has_table', False)),
                "has_equation": bool(getattr(chunk, 'has_equation', False)),
                "chunk_type": str(chunk_type),
                "source": str(getattr(chunk, 'source', '') or ''),
            }

            # Add any additional metadata
            if hasattr(chunk, 'metadata') and isinstance(chunk.metadata, dict):
                for key, value in chunk.metadata.items():
                    if key not in payload:
                        payload[key] = value

            points.append(PointStruct(
                id=point_id,
                vector=vectors,
                payload=payload,
            ))

        # Upsert in batches
        total_added = 0
        for i in range(0, len(points), batch_size):
            batch = points[i:i + batch_size]
            self._client.upsert(
                collection_name=self.config.collection_name,
                points=batch,
                wait=True,
            )
            total_added += len(batch)
            if total_added % 500 == 0:
                logger.info(f"Added {total_added}/{len(points)} chunks...")

        logger.info(f"Added {len(points)} chunks to vector store")
        return point_ids

    async def search(
        self,
        query_embedding: list[float],
        top_k: int = 10,
        filters: dict | None = None,
        score_threshold: float | None = None,
    ) -> list[SearchResult]:
        """
        Search for similar chunks using dense embedding.

        Args:
            query_embedding: Query embedding vector (1536 dimensions)
            top_k: Number of results to return
            filters: Optional filters for metadata fields
                Examples:
                    {"doc_id": "ANSI/ANS-8.1-2014"}
                    {"source": "nureg", "has_table": True}
            score_threshold: Minimum similarity score threshold

        Returns:
            List of SearchResult objects sorted by relevance
        """

        qdrant_filter = self._build_filter(filters)

        # Search using dense vector (using query_points API)
        results = self._client.query_points(
            collection_name=self.config.collection_name,
            query=query_embedding,
            using="dense",
            limit=top_k,
            query_filter=qdrant_filter,
            score_threshold=score_threshold,
            with_payload=True,
        )

        return [self._point_to_result(r) for r in results.points]

    async def hybrid_search(
        self,
        dense_embedding: list[float],
        sparse_embedding: dict,
        top_k: int = 10,
        filters: dict | None = None,
        fusion_type: str = "rrf",
    ) -> list[SearchResult]:
        """
        Perform hybrid search combining dense and sparse vectors.

        Uses Qdrant's built-in fusion (RRF or DBsf) to combine
        results from dense semantic search and sparse keyword search.

        Args:
            dense_embedding: Dense query embedding (1536 dimensions)
            sparse_embedding: Sparse query embedding dict with 'indices' and 'values'
            top_k: Number of results to return
            filters: Optional metadata filters
            fusion_type: Fusion method - "rrf" (Reciprocal Rank Fusion) or "dbsf"
                (Distribution-Based Score Fusion). RRF is recommended.

        Returns:
            List of SearchResult objects with fused scores
        """
        from qdrant_client.models import (
            Fusion,
            FusionQuery,
            Prefetch,
            SparseVector,
        )

        qdrant_filter = self._build_filter(filters)

        # Create prefetch for dense search
        prefetch_dense = Prefetch(
            query=dense_embedding,
            using="dense",
            limit=top_k * 2,  # Over-fetch for better fusion
            filter=qdrant_filter,
        )

        # Create prefetch for sparse search
        sparse_vector = SparseVector(
            indices=sparse_embedding.get("indices", []),
            values=sparse_embedding.get("values", []),
        )

        prefetch_sparse = Prefetch(
            query=sparse_vector,
            using="sparse",
            limit=top_k * 2,
            filter=qdrant_filter,
        )

        # Select fusion method
        fusion_method = Fusion.RRF if fusion_type.lower() == "rrf" else Fusion.DBSf

        # Execute hybrid search with fusion
        results = self._client.query_points(
            collection_name=self.config.collection_name,
            prefetch=[prefetch_dense, prefetch_sparse],
            query=FusionQuery(fusion=fusion_method),
            limit=top_k,
            with_payload=True,
        )

        return [self._point_to_result(r) for r in results.points]

    async def search_by_document(
        self,
        query_embedding: list[float],
        doc_id: str,
        top_k: int = 10,
    ) -> list[SearchResult]:
        """
        Search within a specific document.

        Args:
            query_embedding: Query embedding vector
            doc_id: Human-readable document ID (e.g., "ANSI/ANS-8.1-2014")
            top_k: Number of results

        Returns:
            List of SearchResult objects from the specified document
        """
        return await self.search(
            query_embedding=query_embedding,
            top_k=top_k,
            filters={"doc_id": doc_id},
        )

    async def search_by_source(
        self,
        query_embedding: list[float],
        source: str,
        top_k: int = 10,
    ) -> list[SearchResult]:
        """
        Search within documents from a specific source.

        Args:
            query_embedding: Query embedding vector
            source: Document source (e.g., "ansi", "nureg", "osti")
            top_k: Number of results

        Returns:
            List of SearchResult objects from the specified source
        """
        return await self.search(
            query_embedding=query_embedding,
            top_k=top_k,
            filters={"source": source},
        )

    async def get_chunk_by_id(self, chunk_id: str) -> SearchResult | None:
        """
        Retrieve a specific chunk by ID.

        Args:
            chunk_id: The chunk/point ID

        Returns:
            SearchResult or None if not found
        """
        try:
            points = self._client.retrieve(
                collection_name=self.config.collection_name,
                ids=[chunk_id],
                with_payload=True,
            )
            if points:
                # Create a mock score since retrieve doesn't return scores
                point = points[0]
                point.score = 1.0
                return self._point_to_result(point)
        except Exception as e:
            logger.warning(f"Could not retrieve chunk {chunk_id}: {e}")

        return None

    async def get_chunks_by_document(
        self,
        document_id: str,
        limit: int = 1000,
    ) -> list[SearchResult]:
        """
        Get all chunks for a document.

        Args:
            document_id: Document UUID
            limit: Maximum chunks to return

        Returns:
            List of SearchResult objects
        """
        from qdrant_client.models import FieldCondition, Filter, MatchValue

        results, _ = self._client.scroll(
            collection_name=self.config.collection_name,
            scroll_filter=Filter(
                must=[FieldCondition(
                    key="document_id",
                    match=MatchValue(value=document_id),
                )]
            ),
            limit=limit,
            with_payload=True,
        )

        # Add mock scores
        for r in results:
            r.score = 1.0

        return [self._point_to_result(r) for r in results]

    def _build_filter(self, filters: dict | None):
        """Build Qdrant filter from dict."""
        if not filters:
            return None

        from qdrant_client.models import FieldCondition, Filter, MatchValue, Range

        conditions = []
        for key, value in filters.items():
            if isinstance(value, bool):
                conditions.append(FieldCondition(
                    key=key,
                    match=MatchValue(value=value),
                ))
            elif isinstance(value, (int, float)):
                conditions.append(FieldCondition(
                    key=key,
                    match=MatchValue(value=value),
                ))
            elif isinstance(value, str):
                conditions.append(FieldCondition(
                    key=key,
                    match=MatchValue(value=value),
                ))
            elif isinstance(value, dict):
                # Support range queries: {"page_number": {"gte": 5, "lte": 10}}
                range_params = {}
                if "gte" in value:
                    range_params["gte"] = value["gte"]
                if "gt" in value:
                    range_params["gt"] = value["gt"]
                if "lte" in value:
                    range_params["lte"] = value["lte"]
                if "lt" in value:
                    range_params["lt"] = value["lt"]
                if range_params:
                    conditions.append(FieldCondition(
                        key=key,
                        range=Range(**range_params),
                    ))
            elif isinstance(value, list):
                # Support "any of" queries: {"source": ["ansi", "nureg"]}
                from qdrant_client.models import MatchAny
                conditions.append(FieldCondition(
                    key=key,
                    match=MatchAny(any=value),
                ))

        return Filter(must=conditions) if conditions else None

    def _point_to_result(self, point) -> SearchResult:
        """Convert Qdrant point to SearchResult."""
        payload = point.payload or {}

        return SearchResult(
            chunk_id=payload.get("chunk_id", str(point.id)),
            document_id=payload.get("document_id", ""),
            doc_id=payload.get("doc_id", ""),
            content=payload.get("content", ""),
            score=getattr(point, 'score', 0.0) or 0.0,
            section_path=payload.get("section_path", ""),
            page_number=payload.get("page_number"),
            chunk_index=payload.get("chunk_index", 0),
            has_table=payload.get("has_table", False),
            has_equation=payload.get("has_equation", False),
            source=payload.get("source", ""),
            metadata={
                k: v for k, v in payload.items()
                if k not in [
                    "chunk_id", "document_id", "doc_id", "content",
                    "section_path", "page_number", "chunk_index",
                    "has_table", "has_equation", "source"
                ]
            },
        )

    async def delete_document(self, document_id: str) -> bool:
        """
        Delete all chunks for a document.

        Args:
            document_id: Document UUID

        Returns:
            True if deletion was successful
        """
        from qdrant_client.models import FieldCondition, Filter, MatchValue

        try:
            self._client.delete(
                collection_name=self.config.collection_name,
                points_selector=Filter(
                    must=[FieldCondition(
                        key="document_id",
                        match=MatchValue(value=document_id),
                    )]
                ),
            )
            logger.info(f"Deleted chunks for document {document_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete document {document_id}: {e}")
            return False

    async def delete_by_doc_id(self, doc_id: str) -> bool:
        """
        Delete all chunks for a document by human-readable ID.

        Args:
            doc_id: Human-readable document ID (e.g., "ANSI/ANS-8.1-2014")

        Returns:
            True if deletion was successful
        """
        from qdrant_client.models import FieldCondition, Filter, MatchValue

        try:
            self._client.delete(
                collection_name=self.config.collection_name,
                points_selector=Filter(
                    must=[FieldCondition(
                        key="doc_id",
                        match=MatchValue(value=doc_id),
                    )]
                ),
            )
            logger.info(f"Deleted chunks for doc_id {doc_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete doc_id {doc_id}: {e}")
            return False

    async def get_collection_info(self) -> dict:
        """
        Get information about the collection.

        Returns:
            Dict with collection statistics
        """
        info = self._client.get_collection(self.config.collection_name)
        return {
            "name": self.config.collection_name,
            "points_count": info.points_count,
            "indexed_vectors_count": info.indexed_vectors_count,
            "segments_count": info.segments_count,
            "status": info.status.value if hasattr(info.status, 'value') else str(info.status),
            "optimizer_status": str(info.optimizer_status),
        }

    async def count_documents(self) -> int:
        """Get total number of unique documents in the store."""
        # This is an approximation using scroll
        seen_docs = set()
        offset = None

        while True:
            results, offset = self._client.scroll(
                collection_name=self.config.collection_name,
                limit=1000,
                offset=offset,
                with_payload=["document_id"],
            )

            for point in results:
                if point.payload and point.payload.get("document_id"):
                    seen_docs.add(point.payload["document_id"])

            if offset is None:
                break

        return len(seen_docs)

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

        if self._async_client:
            await self._async_client.close()
            self._async_client = None

        self._initialized = False
        logger.info("Vector store connection closed")


# Alias for backward compatibility
VectorStore = QdrantVectorStore


class HybridRetriever:
    """
    Combines dense (semantic) and sparse (keyword) search.

    Provides a unified interface for retrieval that automatically
    handles embedding generation and search fusion using Qdrant's
    built-in RRF (Reciprocal Rank Fusion).

    Example:
        retriever = HybridRetriever(vector_store, embedding_model, sparse_encoder)

        # Simple search (uses hybrid if sparse encoder available)
        results = await retriever.search("What is the subcritical mass limit?")

        # Search within specific documents
        results = await retriever.search_standards(
            "uranium oxide moderator ratio",
            top_k=5,
        )

        # Search with filters
        results = await retriever.search(
            "criticality accident",
            filters={"source": "nureg"},
        )
    """

    def __init__(
        self,
        vector_store: QdrantVectorStore,
        embedding_model: Any,  # EmbeddingModel from embeddings.py
        sparse_encoder: Any | None = None,  # SparseEncoder from embeddings.py
        default_top_k: int = 10,
        default_fusion: str = "rrf",
    ):
        """
        Initialize hybrid retriever.

        Args:
            vector_store: QdrantVectorStore instance
            embedding_model: EmbeddingModel for dense embeddings
            sparse_encoder: Optional SparseEncoder for sparse embeddings
            default_top_k: Default number of results
            default_fusion: Default fusion method ("rrf" or "dbsf")
        """
        self.vector_store = vector_store
        self.embedding_model = embedding_model
        self.sparse_encoder = sparse_encoder
        self.default_top_k = default_top_k
        self.default_fusion = default_fusion

    async def search(
        self,
        query: str,
        top_k: int | None = None,
        filters: dict | None = None,
        use_hybrid: bool = True,
        score_threshold: float | None = None,
    ) -> list[SearchResult]:
        """
        Search for relevant chunks.

        Automatically uses hybrid search when a sparse encoder is
        available, falling back to dense-only search otherwise.

        Args:
            query: Search query text
            top_k: Number of results
            filters: Optional metadata filters
            use_hybrid: Whether to use hybrid search (if sparse encoder available)
            score_threshold: Minimum score threshold

        Returns:
            List of SearchResult objects sorted by relevance
        """
        top_k = top_k or self.default_top_k

        # Generate dense embedding
        dense_embedding = await self.embedding_model.embed(query)

        # Use hybrid search if sparse encoder available
        if use_hybrid and self.sparse_encoder:
            sparse_embedding = await self.sparse_encoder.encode(query)
            return await self.vector_store.hybrid_search(
                dense_embedding=dense_embedding,
                sparse_embedding=sparse_embedding,
                top_k=top_k,
                filters=filters,
                fusion_type=self.default_fusion,
            )

        # Fall back to dense-only search
        return await self.vector_store.search(
            query_embedding=dense_embedding,
            top_k=top_k,
            filters=filters,
            score_threshold=score_threshold,
        )

    async def search_with_context(
        self,
        query: str,
        context_window: int = 1,
        top_k: int | None = None,
        filters: dict | None = None,
    ) -> list[SearchResult]:
        """
        Search and include surrounding chunks for context.

        Retrieves the matching chunks plus adjacent chunks based
        on chunk_index to provide more context.

        Args:
            query: Search query text
            context_window: Number of chunks before/after to include
            top_k: Number of primary results
            filters: Optional metadata filters

        Returns:
            List of SearchResult objects with context
        """
        top_k = top_k or self.default_top_k

        # Get primary results
        results = await self.search(query, top_k=top_k, filters=filters)

        if not results or context_window == 0:
            return results

        # Collect chunk indices and document IDs
        expanded_results = []
        seen_chunks = set()

        for result in results:
            if result.chunk_id in seen_chunks:
                continue

            seen_chunks.add(result.chunk_id)
            expanded_results.append(result)

            # Get surrounding chunks from the same document
            doc_id = result.document_id
            chunk_idx = result.chunk_index

            for offset in range(-context_window, context_window + 1):
                if offset == 0:
                    continue

                target_idx = chunk_idx + offset
                if target_idx < 0:
                    continue

                # Search for the adjacent chunk
                adjacent = await self.vector_store.search(
                    query_embedding=await self.embedding_model.embed(""),  # Dummy
                    top_k=1,
                    filters={
                        "document_id": doc_id,
                        "chunk_index": target_idx,
                    },
                )

                if adjacent and adjacent[0].chunk_id not in seen_chunks:
                    seen_chunks.add(adjacent[0].chunk_id)
                    expanded_results.append(adjacent[0])

        return expanded_results

    async def search_by_doc_type(
        self,
        query: str,
        doc_type: str,
        top_k: int | None = None,
    ) -> list[SearchResult]:
        """
        Search within a specific document type.

        Args:
            query: Search query text
            doc_type: Document type filter
            top_k: Number of results

        Returns:
            List of SearchResult objects
        """
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"doc_type": doc_type},
        )

    async def search_standards(
        self,
        query: str,
        top_k: int | None = None,
    ) -> list[SearchResult]:
        """
        Search specifically in ANSI/ANS standards documents.

        Args:
            query: Search query text
            top_k: Number of results

        Returns:
            List of SearchResult objects from standards
        """
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"source": "ansi"},
        )

    async def search_nuregs(
        self,
        query: str,
        top_k: int | None = None,
    ) -> list[SearchResult]:
        """
        Search specifically in NUREG documents.

        Args:
            query: Search query text
            top_k: Number of results

        Returns:
            List of SearchResult objects from NUREGs
        """
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"source": "nureg"},
        )

    async def search_tables(
        self,
        query: str,
        top_k: int | None = None,
        filters: dict | None = None,
    ) -> list[SearchResult]:
        """
        Search specifically for table content.

        Args:
            query: Search query text
            top_k: Number of results
            filters: Additional filters

        Returns:
            List of SearchResult objects containing tables
        """
        combined_filters = {"has_table": True}
        if filters:
            combined_filters.update(filters)

        return await self.search(
            query=query,
            top_k=top_k,
            filters=combined_filters,
        )

    async def multi_query_search(
        self,
        queries: list[str],
        top_k: int | None = None,
        filters: dict | None = None,
    ) -> list[SearchResult]:
        """
        Search with multiple query variations and combine results.

        Useful for query expansion where you want to search
        with the original query plus paraphrased versions.

        Args:
            queries: List of query strings (original + variations)
            top_k: Number of results per query
            filters: Optional metadata filters

        Returns:
            De-duplicated list of SearchResult objects
        """
        top_k = top_k or self.default_top_k
        all_results = {}  # chunk_id -> (result, best_score)

        for query in queries:
            results = await self.search(
                query=query,
                top_k=top_k,
                filters=filters,
            )

            for result in results:
                if result.chunk_id not in all_results:
                    all_results[result.chunk_id] = result
                elif result.score > all_results[result.chunk_id].score:
                    all_results[result.chunk_id] = result

        # Sort by score and return top_k
        sorted_results = sorted(
            all_results.values(),
            key=lambda x: x.score,
            reverse=True,
        )

        return sorted_results[:top_k]
