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

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

import toml
from langchain_core.language_models import BaseLLM

from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
from local_deep_research.web_search_engines.search_engine_factory import create_search_engine
from local_deep_research.config import LOCAL_COLLECTIONS_FILE

# Setup logging
logger = logging.getLogger(__name__)

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,
        **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
            **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)
                
        # Find all local collection search engines
        self.local_engines = {}
        try:
            local_collections = toml.load(LOCAL_COLLECTIONS_FILE)

            for collection_id, collection in local_collections.items():
                if not collection.get("enabled", True):
                    continue
                    
                # Create a search engine for this collection
                try:
                    engine = create_search_engine(
                        collection_id,
                        llm=llm,
                        max_filtered_results=max_filtered_results
                    )
                    
                    if engine:
                        self.local_engines[collection_id] = {
                            "engine": engine,
                            "name": collection.get("name", collection_id),
                            "description": collection.get("description", "")
                        }
                except Exception as e:
                    logger.error(f"Error creating search engine for collection '{collection_id}': {e}")
        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 as e:
                logger.error(f"Error searching collection '{collection_id}': {e}")
        
        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 as e:
                logger.error(f"Error getting full content from collection '{collection_id}': {e}")
                # 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
