import os
import concurrent.futures
from typing import Any, Dict, List, Optional
from threading import Lock
import atexit

from loguru import logger

from ...config.search_config import get_setting_from_snapshot
from ...database.thread_local_session import cleanup_current_thread
from ...utilities.enums import SearchMode
from ...utilities.json_utils import extract_json, get_llm_response_text
from ...utilities.thread_context import (
    clear_search_context,
    get_search_context,
    set_search_context,
)
from ...web.services.socket_service import SocketIOService
from ..search_engine_base import BaseSearchEngine
from ..search_engine_factory import create_search_engine
from ..search_engines_config import get_available_engines

# Global thread pool shared by all ParallelSearchEngine instances
# This prevents creating multiple thread pools and having more threads than expected
_global_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None
_global_executor_lock = Lock()


def _get_global_executor(
    max_workers: Optional[int] = None,
) -> Optional[concurrent.futures.ThreadPoolExecutor]:
    """
    Get or initialize the global thread pool executor.
    Thread-safe lazy initialization ensures only one pool is created.

    Args:
        max_workers: Number of worker threads. If None, uses Python's recommended
                    formula: min(32, (os.cpu_count() or 1) + 4) for I/O-bound operations.
                    Only used on first initialization; subsequent calls ignore this parameter.

    Returns:
        The global ThreadPoolExecutor instance, or None if initialization fails
    """
    global _global_executor

    with _global_executor_lock:
        if _global_executor is None:
            if max_workers is None:
                max_workers = min(32, (os.cpu_count() or 1) + 4)

            try:
                _global_executor = concurrent.futures.ThreadPoolExecutor(
                    max_workers=max_workers,
                    thread_name_prefix="parallel_search_",
                )
                logger.info(
                    f"Initialized global ThreadPool with {max_workers} workers "
                    f"(shared by all ParallelSearchEngine instances)"
                )
            except Exception:
                logger.exception(
                    "Failed to create global ThreadPoolExecutor, parallel search will not work"
                )
                return None

        return _global_executor


def shutdown_global_executor(wait: bool = True):
    """
    Shutdown the global thread pool executor.

    This is called automatically at process exit via atexit.
    After calling this, any new ParallelSearchEngine instances will create a new pool.

    Args:
        wait: If True, wait for all threads to complete before returning
    """
    global _global_executor

    with _global_executor_lock:
        if _global_executor is not None:
            try:
                _global_executor.shutdown(wait=wait)
                logger.info("Global ThreadPool shutdown complete")
            except Exception:
                logger.exception("Error shutting down global ThreadPool")
            finally:
                _global_executor = None


# Register automatic cleanup at process exit
atexit.register(lambda: shutdown_global_executor(wait=True))


class ParallelSearchEngine(BaseSearchEngine):
    """
    Parallel search engine that executes multiple search engines simultaneously.
    Uses LLM to select appropriate engines based on query, then runs them all in parallel.

    Thread Pool Management:
        All instances share a single global thread pool to prevent thread proliferation.
        The pool is automatically cleaned up at process exit.

    Usage:
        engine = ParallelSearchEngine(llm)
        results = engine.run("query")
        # No manual cleanup needed - global pool is shared and cleaned up automatically
    """

    def __init__(
        self,
        llm,
        max_results: int = 10,
        use_api_key_services: bool = True,
        max_engines_to_select: int = 100,  # Allow selecting all available engines
        allow_local_engines: bool = False,  # Disabled by default for privacy
        include_generic_engines: bool = True,  # Always include generic search engines
        search_mode: SearchMode = SearchMode.ALL,
        max_filtered_results: Optional[int] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        programmatic_mode: bool = False,
        max_workers: Optional[
            int
        ] = None,  # Thread pool size (None = auto-detect)
        **kwargs,
    ):
        """
        Initialize the parallel search engine.

        All instances share a global thread pool. The first instance created will
        initialize the pool with the specified max_workers (or auto-detected value).
        Subsequent instances reuse the existing pool.

        Args:
            llm: Language model instance for query classification
            max_results: Maximum number of search results to return per engine
            use_api_key_services: Whether to include services that require API keys
            max_engines_to_select: Maximum number of engines to select for parallel execution
            allow_local_engines: Whether to include local/private engines (WARNING: May expose personal data to web)
            include_generic_engines: Always include generic search engines (searxng, brave, ddg, etc.)
            search_mode: SearchMode enum value - ALL for all engines, SCIENTIFIC for scientific + generic engines only
            max_filtered_results: Maximum number of results to keep after filtering
            settings_snapshot: Settings snapshot for thread context
            programmatic_mode: If True, disables database operations and metrics tracking
            max_workers: Thread pool size for the FIRST instance only. If None, uses Python's
                        recommended formula: min(32, (os.cpu_count() or 1) + 4) for I/O-bound operations.
                        Ignored if the global pool is already initialized.
            **kwargs: Additional parameters (ignored but accepted for compatibility)
        """
        # Parallel search aggregates results from multiple engines, so
        # it runs with a higher post-filter cap than a single-engine
        # search would get from DEFAULT_MAX_FILTERED_RESULTS.
        if max_filtered_results is None:
            max_filtered_results = 50
            logger.info(
                f"Setting max_filtered_results to {max_filtered_results} for parallel search"
            )

        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            settings_snapshot=settings_snapshot,
            programmatic_mode=programmatic_mode,
        )

        self.use_api_key_services = use_api_key_services
        self.max_engines_to_select = max_engines_to_select
        self.allow_local_engines = allow_local_engines
        self.include_generic_engines = include_generic_engines
        self.search_mode = search_mode
        self.settings_snapshot = settings_snapshot or {}

        # Disable LLM relevance filtering at the parallel level by default
        # Individual engines can still filter their own results
        # Double filtering (engines + parallel) is too aggressive
        self.enable_llm_relevance_filter = kwargs.get(
            "enable_llm_relevance_filter", False
        )

        # Cache for engine instances
        self.engine_cache: Dict[str, Any] = {}
        self.cache_lock = Lock()

        # Initialize global thread pool (thread-safe, only happens once)
        # All instances share the same pool to prevent thread proliferation
        _get_global_executor(max_workers)

        # Get available engines (excluding 'meta', 'auto', and 'parallel')
        self.available_engines = self._get_available_engines()
        logger.info(
            f"Parallel Search Engine initialized with {len(self.available_engines)} available engines: {', '.join(self.available_engines)}"
        )

    def _get_search_config(self) -> Dict[str, Any]:
        """Get search config for available engines."""
        return get_available_engines(
            settings_snapshot=self.settings_snapshot,
            use_api_key_services=self.use_api_key_services,
        )

    def _get_available_engines(self) -> List[str]:
        """Get list of available engines, applying parallel-specific filters
        (is_public, search mode) on top of the shared base filter."""
        base_available = get_available_engines(
            settings_snapshot=self.settings_snapshot,
            use_api_key_services=self.use_api_key_services,
        )

        available = []
        for name, config_ in base_available.items():
            # Try to get the engine class to check is_public flag
            success, engine_class, error_msg = (
                BaseSearchEngine._load_engine_class(name, config_)
            )

            if not success:
                logger.debug(error_msg)
                continue

            # Check if engine is public or if local engines are allowed
            if hasattr(engine_class, "is_public"):
                if not engine_class.is_public and not self.allow_local_engines:
                    logger.debug(f"Skipping local/private engine: {name}")
                    continue
                if not engine_class.is_public and self.allow_local_engines:
                    logger.warning(
                        f"Including local/private engine {name} - data may be exposed"
                    )
            else:
                # No is_public flag - assume it's private/local for safety
                if not self.allow_local_engines:
                    logger.debug(
                        f"Skipping engine {name} - no is_public flag and local engines not allowed"
                    )
                    continue

            # Apply scientific mode filtering if enabled
            if self.search_mode == SearchMode.SCIENTIFIC:
                is_scientific = getattr(engine_class, "is_scientific", False)
                is_generic = getattr(engine_class, "is_generic", False)

                if not (is_scientific or is_generic):
                    logger.debug(
                        f"Skipping {name} in scientific mode (not scientific or generic)"
                    )
                    continue

                logger.debug(
                    f"Including {name} in scientific mode (scientific={is_scientific}, generic={is_generic})"
                )

            available.append(name)

        return available

    def _get_available_generic_engines(self) -> List[str]:
        """Get list of available generic search engines that pass API key checks"""
        generic_engines = []
        config_data = self._get_search_config()

        for name, config_ in config_data.items():
            # Skip if not in available engines (already filtered for API keys etc)
            if name not in self.available_engines:
                continue

            # Load the engine class to check is_generic flag
            success, engine_class, error_msg = (
                BaseSearchEngine._load_engine_class(name, config_)
            )

            if not success:
                logger.debug(
                    f"Could not check if {name} is generic: {error_msg}"
                )
                continue

            # Check if engine is generic
            if getattr(engine_class, "is_generic", False):
                generic_engines.append(name)
                logger.debug(f"Found generic engine: {name}")

        return generic_engines

    def select_engines(self, query: str) -> List[str]:
        """
        Use LLM to select appropriate search engines based only on names.

        Args:
            query: The search query

        Returns:
            List of selected engine names
        """
        if not self.llm or not self.available_engines:
            logger.warning(
                "No LLM or no available engines, using all available"
            )
            return self.available_engines[: self.max_engines_to_select]

        try:
            # Get list of engines for LLM to select from (exclude generic ones if they'll be auto-added)
            engines_for_selection = self.available_engines.copy()
            generic_engines = []

            if self.include_generic_engines:
                generic_engines = self._get_available_generic_engines()
                # Remove generic engines from selection since they'll be added automatically
                engines_for_selection = [
                    e for e in engines_for_selection if e not in generic_engines
                ]
                logger.debug(
                    f"Excluding generic engines from LLM selection: {generic_engines}"
                )

            # If no specialized engines available, just return the generic ones
            if not engines_for_selection:
                logger.info(
                    f"No specialized engines available, using generic engines: {generic_engines}"
                )
                return generic_engines

            # Create a simple prompt with just non-generic engine names
            engine_list = "\n".join(
                [
                    f"[{i}] {name}"
                    for i, name in enumerate(engines_for_selection)
                ]
            )

            logger.debug(f"Engines for LLM selection: {engines_for_selection}")

            prompt = f"""Query: {query}

Available search engines:
{engine_list}

Select the most appropriate search engines for this query. Return ONLY a JSON array of indices.
Example: [0, 2, 5]

Select up to {self.max_engines_to_select} engines that would best answer this query."""

            logger.debug("Sending prompt to LLM for engine selection")
            # Get LLM response
            response = self.llm.invoke(prompt)
            content = get_llm_response_text(response)

            indices = extract_json(content, expected_type=list)

            if indices is not None:
                # Convert indices to engine names (from the filtered list)
                selected_engines = []
                for idx in indices:
                    if isinstance(idx, int) and 0 <= idx < len(
                        engines_for_selection
                    ):
                        selected_engines.append(engines_for_selection[idx])

                if selected_engines:
                    # Add generic search engines (they were excluded from selection)
                    if self.include_generic_engines:
                        for engine in generic_engines:
                            if engine not in selected_engines:
                                selected_engines.append(engine)
                                logger.debug(f"Added generic engine: {engine}")

                    logger.info(f"Final selected engines: {selected_engines}")
                    return selected_engines

            # Fallback if parsing fails - return generic engines plus some specialized ones
            logger.warning(
                "Failed to parse LLM response, using generic engines + top specialized"
            )
            result = (
                generic_engines.copy() if self.include_generic_engines else []
            )
            for engine in self.available_engines[: self.max_engines_to_select]:
                if engine not in result:
                    result.append(engine)
            return result

        except Exception:
            logger.exception("Error selecting engines with LLM")
            # Fallback to using generic engines + available engines
            if self.include_generic_engines:
                generic_engines = self._get_available_generic_engines()
                result = generic_engines.copy()
                for engine in self.available_engines[
                    : self.max_engines_to_select
                ]:
                    if engine not in result:
                        result.append(engine)
                return result
            return self.available_engines[: self.max_engines_to_select]

    def _get_engine_instance(
        self, engine_name: str
    ) -> Optional[BaseSearchEngine]:
        """Get or create an instance of the specified search engine"""
        with self.cache_lock:
            # Return cached instance if available
            if engine_name in self.engine_cache:
                return self.engine_cache[engine_name]  # type: ignore[no-any-return]

            # Create a new instance
            engine = None
            try:
                # Only pass parameters that all engines accept
                common_params = {
                    "llm": self.llm,
                    "max_results": self.max_results,
                }

                # Add max_filtered_results if specified
                if self.max_filtered_results is not None:
                    common_params["max_filtered_results"] = (
                        self.max_filtered_results
                    )

                engine = create_search_engine(
                    engine_name,
                    settings_snapshot=self.settings_snapshot,
                    programmatic_mode=self.programmatic_mode,
                    **common_params,  # type: ignore[arg-type]
                )

                # Individual engines use their auto-detected filtering settings
                # The factory enables LLM filtering for scientific engines (arXiv, etc.)
                # and disables it for generic engines (Google, Brave, etc.)
            except Exception:
                logger.exception(
                    f"Error creating engine instance for {engine_name}"
                )
                return None

            if engine:
                # Cache the instance
                self.engine_cache[engine_name] = engine

            return engine

    def _execute_single_engine(
        self, engine_name: str, query: str
    ) -> Dict[str, Any]:
        """
        Execute a single search engine and return results.

        Context propagation and thread-local cleanup are handled by
        ``_run_with_context`` at the submit site in ``_get_previews``,
        which captures the submitter's context once per request.

        Args:
            engine_name: Name of the engine to execute
            query: The search query

        Returns:
            Dictionary with engine name and results or error
        """
        logger.info(f"Executing search on {engine_name}")

        engine = self._get_engine_instance(engine_name)
        if not engine:
            return {
                "engine": engine_name,
                "success": False,
                "error": f"Failed to initialize {engine_name}",
                "results": [],
            }

        try:
            # Run the engine properly through its run() method
            # This ensures proper filter application, context propagation, etc.
            results = engine.run(query)

            if results and len(results) > 0:
                logger.info(f"Got {len(results)} results from {engine_name}")
                return {
                    "engine": engine_name,
                    "success": True,
                    "results": results,
                    "count": len(results),
                }
            return {
                "engine": engine_name,
                "success": False,
                "error": "No results",
                "results": [],
            }

        except Exception:
            logger.exception(f"Error executing {engine_name}")
            return {
                "engine": engine_name,
                "success": False,
                "error": f"Engine {engine_name} failed",
                "results": [],
            }

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information by executing selected engines in parallel.

        Args:
            query: The search query

        Returns:
            Combined list of preview dictionaries from all successful engines
        """
        # Select engines for this query
        selected_engines = self.select_engines(query)

        if not selected_engines:
            logger.warning("No engines selected")
            return []

        logger.info(
            f"PARALLEL_SEARCH: Executing {len(selected_engines)} engines in parallel: {', '.join(selected_engines)}"
        )

        # Emit socket event about selected engines
        try:
            SocketIOService().emit_socket_event(
                "parallel_search_started",
                {"engines": selected_engines, "query": query},
            )
        except Exception:
            logger.exception("Socket emit error (non-critical)")

        # Execute all engines in parallel using persistent thread pool
        all_results = []
        engine_results = {}

        # Get the global thread pool
        executor = _get_global_executor()
        if executor is None:
            logger.error(
                "Global thread pool not available, cannot execute parallel search"
            )
            return []

        # Capture submitter's research context ONCE (in the request thread,
        # where it is populated by the strategy's caller). Workers in the
        # global pool do not inherit thread-local state, so we must re-apply
        # it per task and clear it (+ clean up thread-local DB/session/password
        # state) in a finally block so nothing leaks to the next task that
        # lands on the same worker.
        submitter_ctx = get_search_context()

        def _run_with_context(fn, *fargs):
            if submitter_ctx is not None:
                set_search_context(submitter_ctx)
            try:
                return fn(*fargs)
            finally:
                clear_search_context()
                try:
                    cleanup_current_thread()
                except Exception:
                    logger.debug(
                        "parallel_search submit wrapper: cleanup_current_thread failed",
                        exc_info=True,
                    )

        # Submit all tasks to the global executor
        future_to_engine = {
            executor.submit(
                _run_with_context,
                self._execute_single_engine,
                engine_name,
                query,
            ): engine_name
            for engine_name in selected_engines
        }

        # Collect results as they complete
        for future in concurrent.futures.as_completed(future_to_engine):
            engine_name = future_to_engine[future]
            try:
                result = future.result()
                engine_results[engine_name] = result

                if result["success"]:
                    # Add source information to each result
                    for item in result["results"]:
                        item["search_engine"] = engine_name
                    all_results.extend(result["results"])

                    # Emit success event
                    try:
                        SocketIOService().emit_socket_event(
                            "engine_completed",
                            {
                                "engine": engine_name,
                                "success": True,
                                "count": result["count"],
                            },
                        )
                    except Exception:
                        logger.debug("Socket emit error (non-critical)")
                else:
                    # Emit failure event
                    try:
                        SocketIOService().emit_socket_event(
                            "engine_completed",
                            {
                                "engine": engine_name,
                                "success": False,
                                "error": result.get("error", "Unknown error"),
                            },
                        )
                    except Exception:
                        logger.debug("Socket emit error (non-critical)")

            except Exception:
                logger.exception(f"Failed to get result from {engine_name}")
                engine_results[engine_name] = {
                    "engine": engine_name,
                    "success": False,
                    "error": "Search execution failed",
                    "results": [],
                }

        # Log summary
        successful_engines = [
            name for name, res in engine_results.items() if res["success"]
        ]
        failed_engines = [
            name for name, res in engine_results.items() if not res["success"]
        ]

        logger.info(
            f"PARALLEL_SEARCH_COMPLETE: {len(successful_engines)} succeeded, {len(failed_engines)} failed"
        )
        if successful_engines:
            logger.info(f"Successful engines: {', '.join(successful_engines)}")
        if failed_engines:
            logger.info(f"Failed engines: {', '.join(failed_engines)}")

        # Log sample result to understand structure
        if all_results:
            logger.debug(
                f"Sample result keys from first result: {list(all_results[0].keys())}"
            )
            logger.debug(f"Sample result: {str(all_results[0])[:500]}")

        logger.info(f"Total results from all engines: {len(all_results)}")

        # Store the engine results for potential use in _get_full_content
        self._engine_results = engine_results
        self._successful_engines = successful_engines

        return all_results

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant items.
        Since we ran multiple engines, we'll group items by their source engine
        and get full content from each engine for its own results.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        # Check if we should get full content
        if get_setting_from_snapshot(
            "search.snippets_only",
            True,
            settings_snapshot=self.settings_snapshot,
        ):
            logger.info("Snippet-only mode, skipping full content retrieval")
            return relevant_items

        logger.info("Getting full content for relevant items")

        # Group items by their source engine
        items_by_engine: Dict[str, List[Dict[str, Any]]] = {}
        for item in relevant_items:
            engine_name = item.get("search_engine")
            if engine_name:
                if engine_name not in items_by_engine:
                    items_by_engine[engine_name] = []
                items_by_engine[engine_name].append(item)

        # Get full content from each engine for its items
        all_full_content = []

        for engine_name, items in items_by_engine.items():
            engine = self._get_engine_instance(engine_name)
            if engine:
                try:
                    logger.info(
                        f"Getting full content from {engine_name} for {len(items)} items"
                    )
                    full_content = engine._get_full_content(items)
                    all_full_content.extend(full_content)
                except Exception:
                    logger.exception(
                        f"Error getting full content from {engine_name}"
                    )
                    # Fall back to returning items without full content
                    all_full_content.extend(items)
            else:
                # No engine available, return items as-is
                all_full_content.extend(items)

        return all_full_content

    def close(self):
        """Close all cached child search engines and own resources."""
        from ...utilities.resource_utils import safe_close

        for engine in self.engine_cache.values():
            safe_close(engine, "child search engine")
        self.engine_cache.clear()
        super().close()

    def invoke(self, query: str) -> List[Dict[str, Any]]:
        """Compatibility method for LangChain tools"""
        return self.run(query)

    # Note: No shutdown() or context manager methods needed
    # The global thread pool is automatically cleaned up at process exit via atexit
