"""
Semantic search utilities for RapidRAG.
"""

from typing import Optional, List, Dict, Any


class SemanticSearch:
    """
    Semantic search utilities.

    Wraps RapidRAG.search() with additional functionality.
    """

    def __init__(self, rag: "RapidRAG"):
        """Initialize with a RapidRAG instance."""
        self.rag = rag

    def search(
        self,
        query: str,
        n_results: int = 5,
        min_score: float = 0.0,
        source_filter: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Search with additional filters.

        Args:
            query: Search query
            n_results: Max results
            min_score: Minimum similarity score (0-1)
            source_filter: Filter by source path (contains)

        Returns:
            Filtered search results
        """
        results = self.rag.search(query, n_results=n_results * 2)

        # Filter by score
        results = [r for r in results if r["score"] >= min_score]

        # Filter by source
        if source_filter:
            results = [
                r for r in results
                if source_filter in r["metadata"].get("source", "")
            ]

        return results[:n_results]

    def find_similar(
        self,
        doc_id: str,
        n_results: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find documents similar to a given document.

        Args:
            doc_id: ID of the reference document
            n_results: Number of similar documents to find

        Returns:
            List of similar documents
        """
        # Get the document
        result = self.rag.collection.get(ids=[doc_id], include=["documents"])

        if not result["documents"]:
            return []

        content = result["documents"][0]

        # Search for similar (excluding the original)
        results = self.rag.search(content, n_results=n_results + 1)

        # Remove the original document
        return [r for r in results if r["id"] != doc_id][:n_results]

    def hybrid_search(
        self,
        query: str,
        keywords: List[str],
        n_results: int = 5,
        keyword_boost: float = 0.3
    ) -> List[Dict[str, Any]]:
        """
        Hybrid search combining semantic + keyword matching.

        Args:
            query: Semantic search query
            keywords: Keywords to boost
            n_results: Number of results
            keyword_boost: Score boost for keyword matches (0-1)

        Returns:
            Reranked results
        """
        # Semantic search
        results = self.rag.search(query, n_results=n_results * 2)

        # Boost scores for keyword matches
        for r in results:
            content_lower = r["content"].lower()
            matches = sum(1 for kw in keywords if kw.lower() in content_lower)
            if matches > 0:
                boost = min(keyword_boost * matches, keyword_boost * 3)
                r["score"] = min(1.0, r["score"] + boost)
                r["keyword_matches"] = matches

        # Sort by score
        results.sort(key=lambda x: x["score"], reverse=True)

        return results[:n_results]
