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

from langchain_core.language_models import BaseChatModel

from .advanced_search_system.findings.repository import FindingsRepository
from .advanced_search_system.questions.standard_question import (
    StandardQuestionGenerator,
)
from .advanced_search_system.strategies.iterdrag_strategy import IterDRAGStrategy
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.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

logger = logging.getLogger(__name__)


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

    def __init__(
        self,
        strategy_name: str = "source-based",
        include_text_content: bool = True,
        use_cross_engine_filter: bool = True,
        llm: BaseChatModel | None = None,
        search: BaseSearchEngine | None = None,
    ):
        """Initialize the advanced search system.

        Args:
            strategy_name: The name of the search strategy to use ("standard" or "iterdrag")
            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.
        """
        # 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 = get_db_setting("search.iterations", 1)

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

        # if hasattr(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
        result["questions_by_iteration"] = self.questions_by_iteration
        return result
