"""
Smart Decomposition Strategy that chooses between recursive and adaptive approaches.

This strategy analyzes the query type and chooses the most appropriate
decomposition method for optimal results.
"""

from enum import Enum
from typing import Any, Dict, List

from langchain_core.language_models import BaseChatModel
from loguru import logger

from ...utilities.search_utilities import remove_think_tags
from .adaptive_decomposition_strategy import AdaptiveDecompositionStrategy
from .base_strategy import BaseSearchStrategy
from .browsecomp_optimized_strategy import BrowseCompOptimizedStrategy
from .evidence_based_strategy import EvidenceBasedStrategy
from .recursive_decomposition_strategy import RecursiveDecompositionStrategy


class QueryType(Enum):
    """Types of queries that benefit from different strategies."""

    PUZZLE_LIKE = "puzzle_like"  # Queries with specific clues/constraints
    HIERARCHICAL = "hierarchical"  # Multi-part questions with clear structure
    EXPLORATORY = "exploratory"  # Open-ended research questions
    COMPARATIVE = "comparative"  # Comparison between multiple entities
    FACTUAL = "factual"  # Simple factual questions
    CONSTRAINT_BASED = (
        "constraint_based"  # Queries with multiple verifiable constraints
    )


class SmartDecompositionStrategy(BaseSearchStrategy):
    """
    A meta-strategy that intelligently chooses between decomposition approaches.

    It analyzes the query characteristics and selects either:
    - Adaptive decomposition for puzzle-like or constraint-based queries
    - Recursive decomposition for clearly hierarchical queries
    - Direct search for simple factual queries
    """

    def __init__(
        self,
        model: BaseChatModel,
        search: Any,
        all_links_of_system: List[Any],
        settings_snapshot=None,
        **kwargs,
    ):
        """Initialize the smart decomposition strategy.

        Args:
            model: The language model to use
            search: The search engine instance
            all_links_of_system: List to store all encountered links
            settings_snapshot: Settings snapshot for thread context
            **kwargs: Additional parameters for sub-strategies
        """
        super().__init__(all_links_of_system, settings_snapshot)
        self.model = model
        self.search = search
        self.strategy_params = kwargs

    def analyze_topic(self, query: str) -> Dict:
        """Analyze a topic using the most appropriate strategy.

        Args:
            query: The research query to analyze

        Returns:
            Dictionary containing analysis results
        """
        # Determine query type
        query_type = self._classify_query(query)
        logger.info(f"Query classified as: {query_type.value}")

        # Progress callback for UI
        if self.progress_callback:
            self.progress_callback(
                f"Analyzing query type: {query_type.value}",
                5,
                {
                    "phase": "query_classification",
                    "query_type": query_type.value,
                },
            )

        # Choose and execute appropriate strategy
        if query_type == QueryType.PUZZLE_LIKE:
            # Check if it's specifically a BrowseComp-style puzzle
            if self._is_browsecomp_style(query):
                return self._use_browsecomp_strategy(query)
            return self._use_evidence_strategy(query)
        if query_type == QueryType.CONSTRAINT_BASED:
            return self._use_evidence_strategy(query)
        if query_type == QueryType.HIERARCHICAL or query_type in [
            QueryType.COMPARATIVE,
            QueryType.EXPLORATORY,
        ]:
            return self._use_recursive_strategy(query)
        # FACTUAL or unknown
        return self._use_adaptive_strategy(query)

    def _classify_query(self, query: str) -> QueryType:
        """Classify the query type to determine best strategy.

        Args:
            query: The query to classify

        Returns:
            QueryType enum value
        """
        prompt = f"""Analyze this query and classify its type:

Query: {query}

Query Types:
1. PUZZLE_LIKE - Contains specific clues, constraints, or puzzle elements that need to be solved step by step
2. HIERARCHICAL - Multi-part question with clear subtopics that can be addressed independently
3. EXPLORATORY - Open-ended research question seeking broad understanding
4. COMPARATIVE - Requires comparing multiple entities or concepts
5. FACTUAL - Simple factual question with a direct answer
6. CONSTRAINT_BASED - Has multiple specific constraints that need verification (dates, numbers, properties)

Consider:
- Does it have specific constraints or clues to follow?
- Can it be broken into independent subtasks?
- Is it seeking a specific answer or broad understanding?
- Does it require step-by-step verification?

Respond with:
QUERY_TYPE: [type]
REASONING: [brief explanation]
KEY_CHARACTERISTICS: [list main characteristics]
"""

        response = self.model.invoke(prompt)
        content = remove_think_tags(str(response.content))

        # Parse response
        query_type = QueryType.EXPLORATORY  # Default

        lines = content.strip().split("\n")
        for line in lines:
            if line.startswith("QUERY_TYPE:"):
                type_str = line.split(":", 1)[1].strip()
                try:
                    query_type = QueryType[type_str]
                except KeyError:
                    # Try to match by value
                    for q_type in QueryType:
                        if q_type.value in type_str.lower():
                            query_type = q_type
                            break

        logger.info(f"Query classification reasoning: {content}")
        return query_type

    def _use_adaptive_strategy(self, query: str) -> Dict:
        """Use adaptive decomposition strategy.

        Args:
            query: The query to analyze

        Returns:
            Analysis results
        """
        logger.info("Using adaptive decomposition strategy")

        strategy = AdaptiveDecompositionStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            **self.strategy_params,
        )

        if self.progress_callback:
            strategy.set_progress_callback(self.progress_callback)

        return strategy.analyze_topic(query)

    def _is_browsecomp_style(self, query: str) -> bool:
        """Check if query is a BrowseComp-style puzzle.

        Args:
            query: The query to check

        Returns:
            True if it appears to be a BrowseComp puzzle
        """
        browsecomp_indicators = [
            "specific scenic location",
            "body part",
            "fell from",
            "viewpoint",
            "search and rescue",
            "SAR incidents",
            "exact answer",
            "multiple clues",
            "last ice age",
            "times more",
            "between",
            "what is the name",
        ]

        lower_query = query.lower()
        indicator_count = sum(
            1
            for indicator in browsecomp_indicators
            if indicator.lower() in lower_query
        )

        # If 3+ indicators present, likely BrowseComp-style
        return indicator_count >= 3

    def _use_browsecomp_strategy(self, query: str) -> Dict:
        """Use BrowseComp-optimized strategy.

        Args:
            query: The query to analyze

        Returns:
            Analysis results
        """
        logger.info("Using BrowseComp-optimized strategy for puzzle query")

        strategy = BrowseCompOptimizedStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            **self.strategy_params,
        )

        if self.progress_callback:
            strategy.set_progress_callback(self.progress_callback)

        return strategy.analyze_topic(query)

    def _use_evidence_strategy(self, query: str) -> Dict:
        """Use evidence-based strategy.

        Args:
            query: The query to analyze

        Returns:
            Analysis results
        """
        logger.info("Using evidence-based strategy for constraint-based query")

        strategy = EvidenceBasedStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            **self.strategy_params,
        )

        if self.progress_callback:
            strategy.set_progress_callback(self.progress_callback)

        return strategy.analyze_topic(query)

    def _use_recursive_strategy(self, query: str) -> Dict:
        """Use recursive decomposition strategy.

        Args:
            query: The query to analyze

        Returns:
            Analysis results
        """
        logger.info("Using recursive decomposition strategy")

        strategy = RecursiveDecompositionStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            **self.strategy_params,
        )

        if self.progress_callback:
            strategy.set_progress_callback(self.progress_callback)

        return strategy.analyze_topic(query)
