# src/local_deep_research/search_system/search_system.py
from typing import 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.adaptive_decomposition_strategy import (
    AdaptiveDecompositionStrategy,
)
from .advanced_search_system.strategies.browsecomp_optimized_strategy import (
    BrowseCompOptimizedStrategy,
)
from .advanced_search_system.strategies.constrained_search_strategy import (
    ConstrainedSearchStrategy,
)
from .advanced_search_system.strategies.early_stop_constrained_strategy import (
    EarlyStopConstrainedStrategy,
)
from .advanced_search_system.strategies.evidence_based_strategy_v2 import (
    EnhancedEvidenceBasedStrategy,
)
from .advanced_search_system.strategies.iterative_reasoning_strategy import (
    IterativeReasoningStrategy,
)
from .advanced_search_system.strategies.iterdrag_strategy import (
    IterDRAGStrategy,
)
from .advanced_search_system.strategies.parallel_constrained_strategy import (
    ParallelConstrainedStrategy,
)
from .advanced_search_system.strategies.parallel_search_strategy import (
    ParallelSearchStrategy,
)
from .advanced_search_system.strategies.rapid_search_strategy import (
    RapidSearchStrategy,
)
from .advanced_search_system.strategies.recursive_decomposition_strategy import (
    RecursiveDecompositionStrategy,
)
from .advanced_search_system.strategies.smart_decomposition_strategy import (
    SmartDecompositionStrategy,
)
from .advanced_search_system.strategies.source_based_strategy import (
    SourceBasedSearchStrategy,
)
from .advanced_search_system.strategies.standard_strategy import (
    StandardSearchStrategy,
)
from .citation_handler import CitationHandler
from .config.llm_config import get_llm
from .config.search_config import get_search
from .utilities.db_utils import get_db_setting
from .web_search_engines.search_engine_base import BaseSearchEngine


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

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

        Args:
            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
            include_text_content: If False, only includes metadata and links in search results
            use_cross_engine_filter: Whether to filter results across search
                engines.
            llm: LLM to use. If not provided, it will use the default one.
            search: Search engine to use. If not provided, it will use the
                default one.
            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.

        """
        # Get configuration
        self.model = llm
        if llm is None:
            self.model = get_llm()
        self.search = search
        if search is None:
            self.search = get_search(llm_instance=self.model)

        # Get iterations setting
        self.max_iterations = max_iterations
        if self.max_iterations is None:
            self.max_iterations = get_db_setting("search.iterations", 1)
        self.questions_per_iteration = questions_per_iteration
        if self.questions_per_iteration is None:
            self.questions_per_iteration = get_db_setting(
                "search.questions_per_iteration", 3
            )

        # 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)
        self.question_generator = StandardQuestionGenerator(self.model)
        self.findings_repository = FindingsRepository(self.model)
        # For backward compatibility
        self.questions_by_iteration = list()
        self.progress_callback = lambda _1, _2, _3: None
        self.all_links_of_system = list()

        # Initialize strategy based on name
        if strategy_name.lower() == "iterdrag":
            logger.info("Creating IterDRAGStrategy instance")
            self.strategy = IterDRAGStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
            )
        elif strategy_name.lower() == "source-based":
            logger.info("Creating SourceBasedSearchStrategy instance")
            self.strategy = SourceBasedSearchStrategy(
                model=self.model,
                search=self.search,
                include_text_content=include_text_content,
                use_cross_engine_filter=use_cross_engine_filter,
                all_links_of_system=self.all_links_of_system,
                use_atomic_facts=use_atomic_facts,
            )
        elif strategy_name.lower() == "parallel":
            logger.info("Creating ParallelSearchStrategy instance")
            self.strategy = ParallelSearchStrategy(
                model=self.model,
                search=self.search,
                include_text_content=include_text_content,
                use_cross_engine_filter=use_cross_engine_filter,
                all_links_of_system=self.all_links_of_system,
            )
        elif strategy_name.lower() == "rapid":
            logger.info("Creating RapidSearchStrategy instance")
            self.strategy = RapidSearchStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
            )
        elif (
            strategy_name.lower() == "recursive"
            or strategy_name.lower() == "recursive-decomposition"
        ):
            logger.info("Creating RecursiveDecompositionStrategy instance")
            self.strategy = RecursiveDecompositionStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
            )
        elif strategy_name.lower() == "iterative":
            logger.info("Creating IterativeReasoningStrategy instance")
            # For iterative reasoning, use more iterations to allow thorough exploration
            # while keeping search iterations separate
            self.strategy = IterativeReasoningStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,  # Increased reasoning iterations
                confidence_threshold=0.95,  # Increased from 0.85 to require very high confidence
                search_iterations_per_round=self.max_iterations
                or 1,  # Search iterations from settings
                questions_per_search=self.questions_per_iteration,
            )
        elif strategy_name.lower() == "adaptive":
            logger.info("Creating AdaptiveDecompositionStrategy instance")
            self.strategy = AdaptiveDecompositionStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_steps=self.max_iterations,
                min_confidence=0.8,
                source_search_iterations=2,
                source_questions_per_iteration=self.questions_per_iteration,
            )
        elif strategy_name.lower() == "smart":
            logger.info("Creating SmartDecompositionStrategy instance")
            self.strategy = SmartDecompositionStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=self.max_iterations,
                source_search_iterations=2,
                source_questions_per_iteration=self.questions_per_iteration,
            )
        elif strategy_name.lower() == "browsecomp":
            logger.info("Creating BrowseCompOptimizedStrategy instance")
            self.strategy = BrowseCompOptimizedStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_browsecomp_iterations=15,  # BrowseComp strategy main iterations
                confidence_threshold=0.9,
                max_iterations=self.max_iterations,  # Source-based sub-searches
                questions_per_iteration=self.questions_per_iteration,  # Source-based sub-searches
            )
        elif strategy_name.lower() == "evidence":
            logger.info("Creating EnhancedEvidenceBasedStrategy instance")
            self.strategy = EnhancedEvidenceBasedStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,  # Main evidence-gathering iterations
                confidence_threshold=0.95,  # Increased from 0.85 to require very high confidence
                candidate_limit=20,  # Increased from 10
                evidence_threshold=0.9,  # Increased from 0.6 to require strong evidence
                max_search_iterations=self.max_iterations,  # Source-based sub-searches
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_threshold=10,  # Increased from 3
                enable_pattern_learning=True,
            )
        elif strategy_name.lower() == "constrained":
            logger.info("Creating ConstrainedSearchStrategy instance")
            self.strategy = ConstrainedSearchStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,  # Increased from 0.85 to require very high confidence
                candidate_limit=100,  # Increased from 30
                evidence_threshold=0.9,  # Increased from 0.6 to require strong evidence
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,  # Increased from 5
            )
        elif strategy_name.lower() in [
            "parallel-constrained",
            "parallel_constrained",
        ]:
            logger.info("Creating ParallelConstrainedStrategy instance")
            self.strategy = ParallelConstrainedStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,  # Increased from 0.85 to require very high confidence
                candidate_limit=100,
                evidence_threshold=0.9,  # Increased from 0.6 to require strong evidence
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,  # Correct parameter name for parent class
                parallel_workers=100,  # Run up to 100 searches in parallel
            )
        elif strategy_name.lower() in [
            "early-stop-constrained",
            "early_stop_constrained",
        ]:
            logger.info("Creating EarlyStopConstrainedStrategy instance")
            self.strategy = EarlyStopConstrainedStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,
                candidate_limit=100,
                evidence_threshold=0.9,
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,
                parallel_workers=100,  # Increased parallelism as requested
                early_stop_threshold=0.99,  # Stop when we find 99%+ confidence
                concurrent_evaluation=True,  # Evaluate candidates as we find them
            )
        elif strategy_name.lower() in ["smart-query", "smart_query"]:
            from .advanced_search_system.strategies.smart_query_strategy import (
                SmartQueryStrategy,
            )

            logger.info("Creating SmartQueryStrategy instance")
            self.strategy = SmartQueryStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,
                candidate_limit=100,
                evidence_threshold=0.9,
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,
                parallel_workers=100,  # High parallelism
                early_stop_threshold=0.99,  # Stop when we find 99%+ confidence
                concurrent_evaluation=True,  # Evaluate candidates as we find them
                use_llm_query_generation=True,  # Use smart query generation
                queries_per_combination=3,  # Generate multiple queries
            )
        elif strategy_name.lower() in ["dual-confidence", "dual_confidence"]:
            from .advanced_search_system.strategies.dual_confidence_strategy import (
                DualConfidenceStrategy,
            )

            logger.info("Creating DualConfidenceStrategy instance")
            self.strategy = DualConfidenceStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,
                candidate_limit=100,
                evidence_threshold=0.9,
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,
                parallel_workers=100,
                early_stop_threshold=0.95,  # Lower threshold since dual confidence is more accurate
                concurrent_evaluation=True,
                use_llm_query_generation=True,
                queries_per_combination=3,
                use_entity_seeding=True,
                use_direct_property_search=True,
                uncertainty_penalty=0.2,  # Penalty for uncertain evidence
                negative_weight=0.5,  # How much negative evidence counts against
            )
        elif strategy_name.lower() in [
            "dual-confidence-with-rejection",
            "dual_confidence_with_rejection",
        ]:
            from .advanced_search_system.strategies.dual_confidence_with_rejection import (
                DualConfidenceWithRejectionStrategy,
            )

            logger.info("Creating DualConfidenceWithRejectionStrategy instance")
            self.strategy = DualConfidenceWithRejectionStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,
                candidate_limit=100,
                evidence_threshold=0.9,
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,
                parallel_workers=100,
                early_stop_threshold=0.95,
                concurrent_evaluation=True,
                use_llm_query_generation=True,
                queries_per_combination=3,
                use_entity_seeding=True,
                use_direct_property_search=True,
                uncertainty_penalty=0.2,
                negative_weight=0.5,
                rejection_threshold=0.3,  # Reject if negative confidence > 30%
                positive_threshold=0.2,  # Unless positive confidence > 20%
                critical_constraint_rejection=0.2,  # Stricter for critical constraints
            )
        elif strategy_name.lower() in [
            "concurrent-dual-confidence",
            "concurrent_dual_confidence",
        ]:
            from .advanced_search_system.strategies.concurrent_dual_confidence_strategy import (
                ConcurrentDualConfidenceStrategy,
            )

            logger.info("Creating ConcurrentDualConfidenceStrategy instance")
            self.strategy = ConcurrentDualConfidenceStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=20,
                confidence_threshold=0.95,
                candidate_limit=100,
                evidence_threshold=0.9,
                max_search_iterations=self.max_iterations,
                questions_per_iteration=self.questions_per_iteration,
                min_candidates_per_stage=20,
                parallel_workers=10,  # Concurrent evaluation threads
                early_stop_threshold=0.95,
                concurrent_evaluation=True,
                use_llm_query_generation=True,
                queries_per_combination=3,
                use_entity_seeding=True,
                use_direct_property_search=True,
                uncertainty_penalty=0.2,
                negative_weight=0.5,
                rejection_threshold=0.3,
                positive_threshold=0.2,
                critical_constraint_rejection=0.2,
                # Concurrent-specific parameters
                min_good_candidates=3,
                target_candidates=5,
                max_candidates=10,
                min_score_threshold=0.65,
                exceptional_score=0.95,
                quality_plateau_threshold=0.1,
                max_search_time=30.0,
                max_evaluations=30,
            )
        elif strategy_name.lower() in ["modular", "modular-strategy"]:
            from .advanced_search_system.strategies.modular_strategy import (
                ModularStrategy,
            )

            logger.info("Creating ModularStrategy instance")
            self.strategy = ModularStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                constraint_checker_type="dual_confidence",  # dual_confidence, strict, threshold
                exploration_strategy="adaptive",  # parallel, adaptive, constraint_guided, diversity
                early_rejection=True,  # Enable fast pre-screening
                early_stopping=True,  # Enable early stopping
                llm_constraint_processing=True,  # Enable LLM-driven constraint processing
                immediate_evaluation=True,  # Enable immediate candidate evaluation
            )
        elif strategy_name.lower() in ["modular-parallel", "modular_parallel"]:
            from .advanced_search_system.strategies.modular_strategy import (
                ModularStrategy,
            )

            logger.info(
                "Creating ModularStrategy with parallel exploration and dual confidence"
            )
            self.strategy = ModularStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                constraint_checker_type="dual_confidence",  # Concurrent evaluation with +/-/? scoring
                exploration_strategy="parallel",  # Parallel constraint searches
            )
        elif strategy_name.lower() in [
            "focused-iteration",
            "focused_iteration",
        ]:
            from .advanced_search_system.strategies.focused_iteration_strategy import (
                FocusedIterationStrategy,
            )

            logger.info("Creating FocusedIterationStrategy instance")
            # PRESERVE SIMPLEQA PERFORMANCE: Keep proven 96.51% accuracy configuration
            # Original optimal settings: max_iterations=8, questions_per_iteration=5
            self.strategy = FocusedIterationStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
                max_iterations=8,  # PROVEN OPTIMAL FOR SIMPLEQA (96.51% accuracy)
                questions_per_iteration=5,  # PROVEN OPTIMAL FOR SIMPLEQA
                use_browsecomp_optimization=True,  # Enable BrowseComp optimizations
            )
        elif strategy_name.lower() in [
            "browsecomp-entity",
            "browsecomp_entity",
        ]:
            from .advanced_search_system.strategies.browsecomp_entity_strategy import (
                BrowseCompEntityStrategy,
            )

            logger.info("Creating BrowseCompEntityStrategy instance")
            self.strategy = BrowseCompEntityStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
            )
        else:
            logger.info("Creating StandardSearchStrategy instance")
            self.strategy = StandardSearchStrategy(
                model=self.model,
                search=self.search,
                all_links_of_system=self.all_links_of_system,
            )

        # 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:
            self.strategy.set_progress_callback(self.progress_callback)

    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) -> Dict:
        """Analyze a topic using the current strategy.

        Args:
            query: The research query to analyze
        """

        # Send progress message with LLM info
        self.progress_callback(
            f"Using {get_db_setting('llm.provider')} model: {get_db_setting('llm.model')}",
            1,  # Low percentage to show this as an early step
            {
                "phase": "setup",
                "llm_info": {
                    "name": get_db_setting("llm.model"),
                    "provider": get_db_setting("llm.provider"),
                },
            },
        )
        # Send progress message with search strategy info
        search_tool = get_db_setting("search.tool")

        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
        return result
