# src/local_deep_research/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,
)

from .utilities.type_utils import unwrap_setting
from .citation_handler import CitationHandler
from .web_search_engines.search_engine_base import BaseSearchEngine
from .constants import DEFAULT_SEARCH_TOOL


def _ensure_snapshot_username(settings_snapshot, username):
    """Ensure the run's settings snapshot carries the username under the
    ``_username`` key that snapshot-driven consumers look it up by.

    The LangGraph agent re-instantiates the search engine *per tool call*
    from this snapshot (``_make_web_search_tool`` -> ``create_search_engine``),
    and registering a user's document collections needs the username
    (``get_user_db_session``), which collection registration reads only from
    ``settings_snapshot["_username"]``. The snapshot is built without it (the
    user lives under ``system.user``), so without this a collection/library
    primary fails inside the agent with "Unknown search engine 'collection_…'".

    Injected here in ``AdvancedSearchSystem.__init__`` — the narrowest common
    consumer of the strategy-running paths — so the web run and the
    programmatic API are repaired in one place. (No-ops when ``username`` is
    absent, e.g. benchmarks/news, which don't use per-user collections anyway.)
    Mirrors the convention ``get_search()`` already applies to its own local
    snapshot copy.

    Returns the snapshot unchanged when there is no username, it is not a dict,
    or ``_username`` is already set (never overwrites an explicit value);
    otherwise a shallow copy with ``_username`` added (no mutation).
    """
    if (
        username
        and isinstance(settings_snapshot, dict)
        and not settings_snapshot.get("_username")
    ):
        return {**settings_snapshot, "_username": username}
    return settings_snapshot


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. Inject the username under "_username" so
        # snapshot-driven consumers (notably the LangGraph agent's per-call
        # search-engine creation, which must register the user's document
        # collections) can find it. See _ensure_snapshot_username.
        self.settings_snapshot = _ensure_snapshot_username(
            settings_snapshot or {}, username
        )

        # 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:
                self.max_iterations = unwrap_setting(
                    self.settings_snapshot["search.iterations"]
                )
            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:
                self.questions_per_iteration = unwrap_setting(
                    self.settings_snapshot["search.questions_per_iteration"]
                )
            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())

        # Defense-in-depth: arm the PEP-578 egress audit-hook backstop for this
        # run if no caller already armed it. The web worker arms it in
        # research_service before calling us; CLI, the news scheduler and the
        # programmatic API construct AdvancedSearchSystem directly and would
        # otherwise run the full pipeline with the secondary net inactive.
        _armed_egress = self._arm_egress_backstop()
        try:
            # Perform the search
            return self._perform_search(
                query, search_id, is_user_search, is_news_search, user_id
            )
        finally:
            if _armed_egress:
                # Only clear what WE armed, so a reused thread doesn't leak the
                # context to unrelated subsequent work (the web worker owns its
                # own teardown via @thread_cleanup).
                from .security.egress.audit_hook import clear_active_context

                clear_active_context()

    def _arm_egress_backstop(self) -> bool:
        """Arm the egress audit-hook context for this run when no caller has.

        Returns True only when THIS call armed the context (so ``analyze_topic``
        clears it on exit). Returns False when it was already armed (web worker)
        or could not be built. Never raises — a backstop failure must not break
        a research run; the explicit PEPs remain the primary enforcement.
        """
        try:
            from .security.egress.audit_hook import (
                get_active_context,
                set_active_context,
            )
            from .security.egress.policy import (
                PolicyDeniedError,
                context_from_snapshot,
            )
        except Exception:
            return False

        if get_active_context() is not None:
            return False  # already armed (e.g. by the web research worker)
        if not self.settings_snapshot:
            return False

        try:
            primary = unwrap_setting(
                self.settings_snapshot.get("search.tool", DEFAULT_SEARCH_TOOL)
            )
            ctx = context_from_snapshot(
                self.settings_snapshot,
                primary or DEFAULT_SEARCH_TOOL,
                username=self.username,
            )
        except (PolicyDeniedError, ValueError):
            # Corrupted scope: the explicit PEPs and the run-start
            # precheck handle these; don't arm a partial backstop.
            logger.bind(policy_audit=True).debug(
                "egress backstop not armed: policy unevaluable", exc_info=True
            )
            return False
        except Exception:
            return False

        set_active_context(ctx)
        return True

    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", DEFAULT_SEARCH_TOOL)
                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
