"""
Search executor module.
Handles the execution of search queries across multiple search engines.
"""

import os
import json
import time
import asyncio
import concurrent.futures
from typing import Dict, List, Any, Optional, Union

from config.config import get_config
from .api_handlers.base_handler import BaseSearchHandler
from .api_handlers.serper_handler import SerperSearchHandler
from .api_handlers.scholar_handler import ScholarSearchHandler
from .api_handlers.arxiv_handler import ArxivSearchHandler


class SearchExecutor:
    """
    Executes search queries across multiple search engines.
    Manages rate limiting, error handling, and result aggregation.
    """

    def __init__(self):
        """Initialize the search executor with available search handlers."""
        self.config = get_config()
        self.handlers = self._initialize_handlers()
        self.available_handlers = {name: handler for name, handler in self.handlers.items() 
                                  if handler.is_available()}

    def _initialize_handlers(self) -> Dict[str, BaseSearchHandler]:
        """
        Initialize all search handlers.

        Returns:
            Dictionary mapping handler names to handler instances
        """
        return {
            "serper": SerperSearchHandler(),
            "scholar": ScholarSearchHandler(),
            "arxiv": ArxivSearchHandler()
        }

    def get_available_search_engines(self) -> List[str]:
        """
        Get a list of available search engines.

        Returns:
            List of available search engine names
        """
        return list(self.available_handlers.keys())

    def execute_search(self, 
                      structured_query: Dict[str, Any], 
                      search_engines: Optional[List[str]] = None, 
                      num_results: int = 10,
                      timeout: int = 30) -> Dict[str, List[Dict[str, Any]]]:
        """
        Execute a search query across multiple search engines.
        
        Args:
            structured_query: Structured query from the query processor
            search_engines: List of search engines to use (if None, use all available)
            num_results: Number of results to return per search engine
            timeout: Timeout in seconds for each search engine
            
        Returns:
            Dictionary mapping search engine names to lists of search results
        """
        # Get the raw query
        raw_query = structured_query.get("raw_query", "")
        
        # Get the enhanced query if available, otherwise use the raw query
        query = structured_query.get("enhanced_query", raw_query)
        
        # Truncate the query if it's too long (Serper API has a 2048 character limit)
        if len(query) > 2000:
            query = query[:2000]
        
        # If no search engines specified, use all available
        if search_engines is None:
            search_engines = list(self.available_handlers.keys())
        else:
            # Filter to only include available search engines
            search_engines = [engine for engine in search_engines 
                             if engine in self.available_handlers]
        
        # Get the search queries for each engine
        search_queries = structured_query.get("search_queries", {})
        
        # Execute searches in parallel
        results = {}
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future_to_engine = {}
            
            for engine in search_engines:
                if engine not in self.available_handlers:
                    continue
                    
                # Get the appropriate query for this engine
                engine_query = search_queries.get(engine, query)
                
                # Submit the search task
                future = executor.submit(
                    self._execute_single_search,
                    engine=engine,
                    query=engine_query,
                    num_results=num_results
                )
                future_to_engine[future] = engine
            
            # Collect results as they complete
            for future in concurrent.futures.as_completed(future_to_engine, timeout=timeout):
                engine = future_to_engine[future]
                try:
                    engine_results = future.result()
                    results[engine] = engine_results
                except Exception as e:
                    print(f"Error executing search for {engine}: {e}")
                    results[engine] = []
        
        return results

    def _execute_single_search(self, engine: str, query: str, num_results: int) -> List[Dict[str, Any]]:
        """
        Execute a search on a single search engine.

        Args:
            engine: Name of the search engine
            query: Query to execute
            num_results: Number of results to return

        Returns:
            List of search results
        """
        handler = self.available_handlers.get(engine)
        if not handler:
            return []
            
        try:
            # Execute the search
            results = handler.search(query, num_results=num_results)
            return results
        except Exception as e:
            print(f"Error executing search for {engine}: {e}")
            return []

    async def execute_search_async(self, 
                                 structured_query: Dict[str, Any], 
                                 search_engines: Optional[List[str]] = None, 
                                 num_results: int = 10,
                                 timeout: int = 30) -> Dict[str, List[Dict[str, Any]]]:
        """
        Execute a search query across specified search engines asynchronously.

        Args:
            structured_query: The structured query from the query processor
            search_engines: List of search engines to use (if None, use all available)
            num_results: Number of results to return per search engine
            timeout: Timeout in seconds for each search engine

        Returns:
            Dictionary mapping search engine names to lists of search results
        """
        # If no search engines specified, use all available
        if search_engines is None:
            search_engines = list(self.available_handlers.keys())
        else:
            # Filter to only include available search engines
            search_engines = [engine for engine in search_engines 
                             if engine in self.available_handlers]
        
        # Get the search queries for each engine
        search_queries = structured_query.get("search_queries", {})
        
        # Create tasks for each search engine
        tasks = []
        for engine in search_engines:
            if engine not in self.available_handlers:
                continue
                
            # Get the appropriate query for this engine
            query = search_queries.get(engine, structured_query.get("enhanced_query", ""))
            
            # Create a task for this search
            task = self._execute_single_search_async(engine, query, num_results)
            tasks.append((engine, task))
        
        # Execute all tasks with timeout
        results = {}
        for engine, task in tasks:
            try:
                engine_results = await asyncio.wait_for(task, timeout=timeout)
                results[engine] = engine_results
            except asyncio.TimeoutError:
                print(f"Search timed out for {engine}")
                results[engine] = []
            except Exception as e:
                print(f"Error executing search for {engine}: {e}")
                results[engine] = []
        
        return results

    async def _execute_single_search_async(self, engine: str, query: str, num_results: int) -> List[Dict[str, Any]]:
        """
        Execute a search on a single search engine asynchronously.

        Args:
            engine: Name of the search engine
            query: Query to execute
            num_results: Number of results to return

        Returns:
            List of search results
        """
        # Execute in a thread pool since most API calls are blocking
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None, self._execute_single_search, engine, query, num_results
        )
