# src/local_deep_research/search_system/search_system.py
from typing import Any, Callable, Dict

from langchain_core.language_models import BaseChatModel
from loguru import logger

from .advanced_search_system.findings.repository import FindingsRepository
from .advanced_search_system.questions.standard_question import (
    StandardQuestionGenerator,
)
from .advanced_search_system.strategies.followup.enhanced_contextual_followup import (
    EnhancedContextualFollowUpStrategy,
)

# StandardSearchStrategy imported lazily to avoid database access during module import
from .citation_handler import CitationHandler
from .web_search_engines.search_engine_base import BaseSearchEngine


class AdvancedSearchSystem:
    """
    Advanced search system that coordinates different search strategies.
    """

    def __init__(
        self,
        llm: BaseChatModel,
        search: BaseSearchEngine,
        strategy_name: str = "source-based",  # Default to comprehensive research strategy
        include_text_content: bool = True,
        use_cross_engine_filter: bool = True,
        max_iterations: int | None = None,
        questions_per_iteration: int | None = None,
        use_atomic_facts: bool = False,
        username: str | None = None,
        settings_snapshot: dict | None = None,
        research_id: str | None = None,
        research_context: dict | None = None,
        programmatic_mode: bool = False,
        search_original_query: bool = True,
    ):
        """Initialize the advanced search system.

        Args:
            llm: LLM to use for the search strategy.
            search: Search engine to use for queries.
            strategy_name: The name of the search strategy to use. Options:
                - "standard": Basic iterative search strategy
                - "iterdrag": Iterative Dense Retrieval Augmented Generation
                - "source-based": Focuses on finding and extracting from sources
                - "parallel": Runs multiple search queries in parallel
                - "rapid": Quick single-pass search
                - "recursive": Recursive decomposition of complex queries
                - "iterative": Loop-based reasoning with persistent knowledge
                - "adaptive": Adaptive step-by-step reasoning
                - "smart": Automatically chooses best strategy based on query
                - "browsecomp": Optimized for BrowseComp-style puzzle queries
                - "evidence": Enhanced evidence-based verification with improved candidate discovery
                - "constrained": Progressive constraint-based search that narrows candidates step by step
                - "parallel-constrained": Parallel constraint-based search with combined constraint execution
                - "early-stop-constrained": Parallel constraint search with immediate evaluation and early stopping at 99% confidence
                - "dual-confidence": Dual confidence scoring with positive/negative/uncertainty
                - "dual-confidence-with-rejection": Dual confidence with early rejection of poor candidates
                - "concurrent-dual-confidence": Concurrent search & evaluation with progressive constraint relaxation
                - "modular": Modular architecture using constraint checking and candidate exploration modules
                - "browsecomp-entity": Entity-focused search for BrowseComp questions with knowledge graph building
                - "iterative-refinement": Iteratively refines results using LLM evaluation and follow-up queries
            include_text_content: If False, only includes metadata and links in search results
            use_cross_engine_filter: Whether to filter results across search
                engines.
            max_iterations: The maximum number of search iterations to
                perform. Will be read from the settings if not specified.
            questions_per_iteration: The number of questions to include in
                each iteration. Will be read from the settings if not specified.
            use_atomic_facts: Whether to use atomic fact decomposition for
                complex queries when using the source-based strategy.
            programmatic_mode: If True, disables database operations and metrics tracking.
                This is useful for running searches without database dependencies.
            search_original_query: Whether to include the original query in the first iteration
                of search. Set to False for news searches to avoid sending long subscription
                prompts to search engines.

        """
        # Store research context for strategies
        self.research_id = research_id
        self.research_context = research_context
        self.username = username

        # Store required components
        self.model = llm
        self.search = search

        # Store settings snapshot
        self.settings_snapshot = settings_snapshot or {}

        # Store programmatic mode
        self.programmatic_mode = programmatic_mode

        # Store search original query setting
        self.search_original_query = search_original_query

        # Log if running in programmatic mode
        if self.programmatic_mode:
            logger.warning(
                "Running in programmatic mode - database operations and metrics tracking disabled. "
                "Rate limiting, search metrics, and persistence features will not be available."
            )

        # Get iterations setting
        self.max_iterations = max_iterations
        if self.max_iterations is None:
            # Use settings from snapshot
            if "search.iterations" in self.settings_snapshot:
                value = self.settings_snapshot["search.iterations"]
                if isinstance(value, dict) and "value" in value:
                    self.max_iterations = value["value"]
                else:
                    self.max_iterations = value
            else:
                self.max_iterations = 1  # Default

        self.questions_per_iteration = questions_per_iteration
        if self.questions_per_iteration is None:
            # Use settings from snapshot
            if "search.questions_per_iteration" in self.settings_snapshot:
                value = self.settings_snapshot["search.questions_per_iteration"]
                if isinstance(value, dict) and "value" in value:
                    self.questions_per_iteration = value["value"]
                else:
                    self.questions_per_iteration = value
            else:
                self.questions_per_iteration = 3  # Default

        # Log the strategy name that's being used
        logger.info(
            f"Initializing AdvancedSearchSystem with strategy_name='{strategy_name}'"
        )

        # Initialize components
        self.citation_handler = CitationHandler(
            self.model, settings_snapshot=self.settings_snapshot
        )
        self.question_generator = StandardQuestionGenerator(self.model)
        self.findings_repository = FindingsRepository(self.model)
        # For backward compatibility
        self.questions_by_iteration: dict[Any, Any] = {}
        self.progress_callback = lambda _1, _2, _3: None
        self.all_links_of_system: list[dict[Any, Any]] = []

        # Initialize strategy using factory
        from .search_system_factory import create_strategy

        # Special handling for follow-up strategy which needs different logic
        if strategy_name.lower() in [
            "enhanced-contextual-followup",
            "enhanced_contextual_followup",
            "contextual-followup",
            "contextual_followup",
        ]:
            logger.info("Creating EnhancedContextualFollowUpStrategy instance")
            # Get delegate strategy from research context
            # This should be the user's preferred strategy from settings
            delegate_strategy_name = (
                self.research_context.get("delegate_strategy", "source-based")
                if self.research_context
                else "source-based"
            )

            delegate = create_strategy(
                strategy_name=delegate_strategy_name,
                model=self.model,
                search=self.search,
                all_links_of_system=[],
                settings_snapshot=self.settings_snapshot,
                knowledge_accumulation_mode=True,
                search_original_query=self.search_original_query,
            )

            # Create the contextual follow-up strategy with the delegate
            self.strategy = EnhancedContextualFollowUpStrategy(
                model=self.model,
                search=self.search,
                delegate_strategy=delegate,
                all_links_of_system=self.all_links_of_system,
                settings_snapshot=self.settings_snapshot,
                research_context=self.research_context,
            )
        else:
            # Use factory for all other strategies
            logger.info(f"Creating {strategy_name} strategy using factory")
            self.strategy = create_strategy(
                strategy_name=strategy_name,
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                settings_snapshot=self.settings_snapshot,
                # Pass strategy-specific parameters
                include_text_content=include_text_content,
                use_cross_engine_filter=use_cross_engine_filter,
                use_atomic_facts=use_atomic_facts,
                max_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                # Special parameters for iterative strategy
                search_iterations_per_round=self.max_iterations or 1,
                questions_per_search=self.questions_per_iteration,
                # Special parameters for adaptive strategy
                max_steps=self.max_iterations,
                source_questions_per_iteration=self.questions_per_iteration,
                # Special parameters for evidence and constrained strategies
                max_search_iterations=self.max_iterations,
                # Special parameters for focused iteration
                use_browsecomp_optimization=True,
                # Pass search original query parameter
                search_original_query=self.search_original_query,
                # Forwarded so strategies that create engines per tool call
                # (e.g. langgraph-agent) can match the system's mode.
                programmatic_mode=self.programmatic_mode,
            )

        # Log the actual strategy class
        logger.info(f"Created strategy of type: {type(self.strategy).__name__}")

        # Configure the strategy with our attributes
        if (
            hasattr(self, "progress_callback")
            and self.progress_callback is not None
        ):
            self.strategy.set_progress_callback(self.progress_callback)

    def close(self):
        """Close resources held by the search system.

        Cascades close to the strategy, which may hold persistent
        ThreadPoolExecutor instances (e.g. ConstraintParallelStrategy holds
        search_executor and evaluation_executor,
        ConcurrentDualConfidenceStrategy holds evaluation_executor).

        NOTE: Does NOT close self.search (the search engine) — the caller
        (run_research_process) manages search engine lifecycle separately
        because the search engine may be shared or reused.

        Most strategies (including the default source-based) use
        context-managed ThreadPoolExecutors that clean up automatically.
        The close() call here is a safety net for the two constraint-based
        strategies that hold persistent executors in __init__. Those
        strategies also shut down their executors in
        find_relevant_information()'s finally block, so this is a second
        line of defense for the edge case where the method is never called.
        """
        from .utilities.resource_utils import safe_close

        if hasattr(self, "strategy"):
            safe_close(self.strategy, "search strategy")

    def _progress_callback(
        self, message: str, progress: int, metadata: dict
    ) -> None:
        """Handle progress updates from the strategy."""
        logger.info(f"Progress: {progress}% - {message}")
        if hasattr(self, "progress_callback"):
            self.progress_callback(message, progress, metadata)

    def set_progress_callback(
        self, callback: Callable[[str, int, dict], None]
    ) -> None:
        """Set a callback function to receive progress updates."""
        self.progress_callback = callback
        if hasattr(self, "strategy"):
            self.strategy.set_progress_callback(callback)

    def analyze_topic(
        self,
        query: str,
        is_user_search: bool = True,
        is_news_search: bool = False,
        user_id: str = "anonymous",
        search_id: str | None = None,
        **kwargs,
    ) -> Dict:
        """Analyze a topic using the current strategy.

        Args:
            query: The research query to analyze
            is_user_search: Whether this is a user-initiated search
            is_news_search: Whether this is a news search
            user_id: The user ID for tracking
            search_id: The search ID (auto-generated if not provided)
            **kwargs: Additional arguments
        """

        # Generate search ID if not provided
        if search_id is None:
            import uuid

            search_id = str(uuid.uuid4())

        # Perform the search
        return self._perform_search(
            query, search_id, is_user_search, is_news_search, user_id
        )

    def _perform_search(
        self,
        query: str,
        search_id: str,
        is_user_search: bool,
        is_news_search: bool,
        user_id: str,
    ) -> Dict:
        """Perform the actual search."""
        # Send progress message with LLM info
        # Get settings from snapshot if available
        llm_provider = "unknown"
        llm_model = "unknown"
        search_tool = "unknown"

        if self.settings_snapshot:
            # Extract values from settings snapshot
            provider_setting = self.settings_snapshot.get("llm.provider", {})
            llm_provider = (
                provider_setting.get("value", "unknown")
                if isinstance(provider_setting, dict)
                else provider_setting
            )

            model_setting = self.settings_snapshot.get("llm.model", {})
            llm_model = (
                model_setting.get("value", "unknown")
                if isinstance(model_setting, dict)
                else model_setting
            )

            tool_setting = self.settings_snapshot.get("search.tool", {})
            search_tool = (
                tool_setting.get("value", "searxng")
                if isinstance(tool_setting, dict)
                else tool_setting
            )

        self.progress_callback(
            f"Using {llm_provider} model: {llm_model}",
            1,  # Low percentage to show this as an early step
            {
                "phase": "setup",
                "llm_info": {
                    "name": llm_model,
                    "provider": llm_provider,
                },
            },
        )
        # Send progress message with search strategy info
        self.progress_callback(
            f"Using search tool: {search_tool}",
            1.5,  # Between setup and processing steps
            {
                "phase": "setup",
                "search_info": {
                    "tool": search_tool,
                },
            },
        )

        # Use the strategy to analyze the topic
        result = self.strategy.analyze_topic(query)

        # Update our attributes for backward compatibility

        self.questions_by_iteration = (
            self.strategy.questions_by_iteration.copy()
        )
        # Send progress message with search info

        # Only extend if they're different objects in memory to avoid duplication
        # This check prevents doubling the list when they reference the same object
        # Fix for issue #301: "too many links in detailed report mode"
        if id(self.all_links_of_system) != id(
            self.strategy.all_links_of_system
        ):
            self.all_links_of_system.extend(self.strategy.all_links_of_system)

        # Include the search system instance for access to citations
        result["search_system"] = self
        result["all_links_of_system"] = self.all_links_of_system

        # Ensure query is included in the result
        if "query" not in result:
            result["query"] = query
        result["questions_by_iteration"] = self.questions_by_iteration

        # Call news callback
        try:
            from .news.core.search_integration import NewsSearchCallback

            callback = NewsSearchCallback()
            context = {
                "is_user_search": is_user_search,
                "is_news_search": is_news_search,
                "user_id": user_id,
                "search_id": search_id,
            }
            callback(query, result, context)
        except Exception:
            logger.exception("Error in news callback")

        return result
