"""
Search engine that searches across all local collections
"""

from typing import Any, Dict, List, Optional, cast

from langchain_core.language_models import BaseLLM
from loguru import logger

from ..search_engine_base import BaseSearchEngine
from ..search_engine_factory import create_search_engine
from ..search_engines_config import local_search_engines
from .search_engine_local import LocalSearchEngine


class LocalAllSearchEngine(BaseSearchEngine):
    """
    Search engine that searches across all local document collections.
    Acts as a meta search engine specifically for local collections.
    """

    def __init__(
        self,
        llm: Optional[BaseLLM] = None,
        max_results: int = 10,
        max_filtered_results: Optional[int] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        programmatic_mode: bool = False,
        **kwargs,
    ):
        """
        Initialize the local all-collections search engine.

        Args:
            llm: Language model for relevance filtering
            max_results: Maximum number of search results
            max_filtered_results: Maximum results after filtering
            settings_snapshot: Settings snapshot for thread context
            programmatic_mode: If True, disables database operations and metrics tracking
            **kwargs: Additional parameters passed to LocalSearchEngine instances
        """
        # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            settings_snapshot=settings_snapshot,
            programmatic_mode=programmatic_mode,
        )

        # Find all local collection search engines
        self.local_engines = {}
        try:
            for collection_id in local_search_engines():
                # Create a search engine for this collection
                try:
                    engine = create_search_engine(
                        collection_id,
                        llm=llm,
                        max_filtered_results=max_filtered_results,
                        settings_snapshot=settings_snapshot,
                        programmatic_mode=programmatic_mode,
                    )
                    engine = cast(LocalSearchEngine, engine)

                    if engine:
                        self.local_engines[collection_id] = {
                            "engine": engine,
                            "name": engine.name,
                            "description": engine.description,
                        }
                except Exception:
                    logger.exception(
                        f"Error creating search engine for collection '{collection_id}'"
                    )
        except ImportError:
            logger.warning("No local collections configuration found")

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information for documents from all local collections.

        Args:
            query: The search query

        Returns:
            List of preview dictionaries
        """
        logger.info(
            f"Searching across all local collections for query: {query}"
        )

        all_previews = []

        # Get previews from each local search engine
        for collection_id, engine_info in self.local_engines.items():
            engine = engine_info["engine"]
            try:
                # Get previews from this engine
                previews = engine._get_previews(query)

                # Add collection info to each preview
                for preview in previews:
                    preview["collection_id"] = collection_id
                    preview["collection_name"] = engine_info["name"]
                    preview["collection_description"] = engine_info[
                        "description"
                    ]

                all_previews.extend(previews)
            except Exception:
                logger.exception(
                    f"Error searching collection '{collection_id}'"
                )

        if not all_previews:
            logger.info(f"No local documents found for query: {query}")
            return []

        # Sort by similarity score if available
        all_previews.sort(
            key=lambda x: float(x.get("similarity", 0)), reverse=True
        )

        # Limit to max_results
        return all_previews[: self.max_results]

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant documents.
        Delegates to the appropriate collection's search engine.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        # Group items by collection
        items_by_collection = {}
        for item in relevant_items:
            collection_id = item.get("collection_id")
            if collection_id and collection_id in self.local_engines:
                if collection_id not in items_by_collection:
                    items_by_collection[collection_id] = []
                items_by_collection[collection_id].append(item)

        # Process each collection's items with its own engine
        all_results = []
        for collection_id, items in items_by_collection.items():
            engine = self.local_engines[collection_id]["engine"]
            try:
                results = engine._get_full_content(items)
                all_results.extend(results)
            except Exception:
                logger.exception(
                    f"Error getting full content from collection '{collection_id}'"
                )
                # Fall back to returning the items without full content
                all_results.extend(items)

        # Add any items that weren't processed
        processed_ids = set(item["id"] for item in all_results)
        for item in relevant_items:
            if item["id"] not in processed_ids:
                all_results.append(item)

        return all_results
