"""
Search executor module.
Handles the execution of search queries across multiple search engines,
including processing of decomposed sub-questions.
"""

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

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

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
from .api_handlers.news_handler import NewsSearchHandler
from .api_handlers.openalex_handler import OpenAlexSearchHandler
from .api_handlers.core_handler import CoreSearchHandler
from .api_handlers.github_handler import GitHubSearchHandler
from .api_handlers.stackexchange_handler import StackExchangeSearchHandler
from .result_enrichers.unpaywall_enricher import UnpaywallEnricher


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()}
        
        # Initialize result enrichers
        self.unpaywall_enricher = UnpaywallEnricher()

    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(),
            "news": NewsSearchHandler(),
            "openalex": OpenAlexSearchHandler(),
            "core": CoreSearchHandler(),
            "github": GitHubSearchHandler(),
            "stackexchange": StackExchangeSearchHandler()
        }

    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())
            
            # Handle specialized query types
            
            # Current events queries
            if structured_query.get("is_current_events", False) and "news" in self.available_handlers:
                print("Current events query detected, prioritizing news search")
                # Make sure news is in the search engines
                if "news" not in search_engines:
                    search_engines.append("news")
                
                # If a specific engine is requested, honor that - otherwise limit to news + a general search engine
                # for a faster response with more relevant results
                if not structured_query.get("specific_engines", False):
                    general_engines = ["serper", "google"]
                    # Find an available general engine
                    general_engine = next((e for e in general_engines if e in self.available_handlers), None)
                    if general_engine:
                        search_engines = ["news", general_engine]
                    else:
                        # Fall back to just news
                        search_engines = ["news"]
            
            # Academic queries
            elif structured_query.get("is_academic", False):
                print("Academic query detected, prioritizing academic search engines")
                
                # Define academic search engines in order of priority
                academic_engines = ["openalex", "core", "arxiv", "scholar"]
                available_academic = [engine for engine in academic_engines if engine in self.available_handlers]
                
                # Always include at least one general search engine for backup
                general_engines = ["serper", "google"]
                available_general = [engine for engine in general_engines if engine in self.available_handlers]
                
                if available_academic and not structured_query.get("specific_engines", False):
                    # Use available academic engines plus one general engine if available
                    search_engines = available_academic
                    if available_general:
                        search_engines.append(available_general[0])
                elif not available_academic:
                    # Just use general search if no academic engines are available
                    search_engines = available_general
                
                print(f"Selected engines for academic query: {search_engines}")
                
            # Code/programming queries
            elif structured_query.get("is_code", False):
                print("Code/programming query detected, prioritizing code search engines")
                
                # Define code search engines in order of priority
                code_engines = ["github", "stackexchange"]
                available_code = [engine for engine in code_engines if engine in self.available_handlers]
                
                # Always include at least one general search engine for backup
                general_engines = ["serper", "google"]
                available_general = [engine for engine in general_engines if engine in self.available_handlers]
                
                if available_code and not structured_query.get("specific_engines", False):
                    # Use available code engines plus one general engine if available
                    search_engines = available_code
                    if available_general:
                        search_engines.append(available_general[0])
                elif not available_code:
                    # Just use general search if no code engines are available
                    search_engines = available_general
                
                print(f"Selected engines for code query: {search_engines}")
        else:
            # Filter to only include available search engines
            search_engines = [engine for engine in search_engines 
                             if engine in self.available_handlers]
            
            # Add specialized handlers based on query type
            
            # For current events queries
            if structured_query.get("is_current_events", False) and "news" in self.available_handlers and "news" not in search_engines:
                print("Current events query detected, adding news search")
                search_engines.append("news")
            
            # For academic queries
            elif structured_query.get("is_academic", False):
                academic_engines = ["openalex", "core", "arxiv", "scholar"]
                for engine in academic_engines:
                    if engine in self.available_handlers and engine not in search_engines:
                        print(f"Academic query detected, adding {engine} search")
                        search_engines.append(engine)
                        
            # For code/programming queries
            elif structured_query.get("is_code", False):
                code_engines = ["github", "stackexchange"]
                for engine in code_engines:
                    if engine in self.available_handlers and engine not in search_engines:
                        print(f"Code query detected, adding {engine} search")
                        search_engines.append(engine)
        
        # Get the search queries for each engine
        search_queries = structured_query.get("search_queries", {})
        
        # For news searches on current events queries, add special parameters
        news_params = {}
        if "news" in search_engines and structured_query.get("is_current_events", False):
            # Set up news search parameters
            news_params["days_back"] = 7  # Limit to 7 days for current events
            news_params["sort_by"] = "publishedAt"  # Sort by publication date
        
        # 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)
                
                # Additional parameters for certain engines
                kwargs = {}
                if engine == "news" and news_params:
                    kwargs = news_params
                
                # Submit the search task
                future = executor.submit(
                    self._execute_single_search,
                    engine=engine,
                    query=engine_query,
                    num_results=num_results,
                    **kwargs
                )
                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, **kwargs) -> 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
            **kwargs: Additional parameters to pass to the search handler

        Returns:
            List of search results
        """
        handler = self.available_handlers.get(engine)
        if not handler:
            return []
            
        try:
            # Execute the search with any additional parameters
            results = handler.search(query, num_results=num_results, **kwargs)
            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
        """
        # Get the enhanced query
        query = structured_query.get("enhanced_query", structured_query.get("original_query", ""))
        
        # If no search engines specified, use all available
        if search_engines is None:
            search_engines = list(self.available_handlers.keys())
            
            # If this is a current events query, prioritize news handler if available
            if structured_query.get("is_current_events", False) and "news" in self.available_handlers:
                print("Current events query detected, prioritizing news search (async)")
                # Make sure news is in the search engines
                if "news" not in search_engines:
                    search_engines.append("news")
                
                # If a specific engine is requested, honor that - otherwise limit to news + a general search engine
                # for a faster response with more relevant results
                if not structured_query.get("specific_engines", False):
                    general_engines = ["serper", "google"]
                    # Find an available general engine
                    general_engine = next((e for e in general_engines if e in self.available_handlers), None)
                    if general_engine:
                        search_engines = ["news", general_engine]
                    else:
                        # Fall back to just news
                        search_engines = ["news"]
        else:
            # Filter to only include available search engines
            search_engines = [engine for engine in search_engines 
                             if engine in self.available_handlers]
            
            # If this is a current events query, add news handler if available and not already included
            if structured_query.get("is_current_events", False) and "news" in self.available_handlers and "news" not in search_engines:
                print("Current events query detected, adding news search (async)")
                search_engines.append("news")
        
        # Get the search queries for each engine
        search_queries = structured_query.get("search_queries", {})
        
        # For news searches on current events queries, add special parameters
        news_params = {}
        if "news" in search_engines and structured_query.get("is_current_events", False):
            # Set up news search parameters
            news_params["days_back"] = 7  # Limit to 7 days for current events
            news_params["sort_by"] = "publishedAt"  # Sort by publication date
        
        # 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
            engine_query = search_queries.get(engine, query)
            
            # Additional parameters for certain engines
            kwargs = {}
            if engine == "news" and news_params:
                kwargs = news_params
            
            # Create a task for this search
            task = self._execute_single_search_async(engine, engine_query, num_results, **kwargs)
            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, **kwargs) -> 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
            **kwargs: Additional parameters to pass to the search handler

        Returns:
            List of search results
        """
        # Execute in a thread pool since most API calls are blocking
        loop = asyncio.get_event_loop()
        
        # Create a partial function with all the arguments
        def execute_search():
            return self._execute_single_search(engine, query, num_results, **kwargs)
        
        return await loop.run_in_executor(None, execute_search)
