"""
Enhanced Contextual Follow-Up Strategy

An improved version of the contextual follow-up strategy that better leverages
past research context, reformulates questions, and reuses sources effectively.
"""

from typing import Dict, List, Optional, Any
from loguru import logger

from ..base_strategy import BaseSearchStrategy
from ...filters.followup_relevance_filter import FollowUpRelevanceFilter
from ...knowledge.followup_context_manager import FollowUpContextHandler
from ...questions.followup.simple_followup_question import (
    SimpleFollowUpQuestionGenerator,
)


class EnhancedContextualFollowUpStrategy(BaseSearchStrategy):
    """
    Enhanced strategy for follow-up research that intelligently uses past context.

    This strategy:
    1. Reformulates follow-up questions based on past findings
    2. Filters and reuses relevant sources from previous research
    3. Passes complete context to the delegate strategy
    4. Optimizes search to avoid redundancy
    """

    def __init__(
        self,
        model,
        search,
        delegate_strategy: BaseSearchStrategy,
        all_links_of_system=None,
        settings_snapshot=None,
        research_context: Optional[Dict] = None,
        **kwargs,
    ):
        """
        Initialize the enhanced contextual follow-up strategy.

        Args:
            model: The LLM model to use
            search: The search engine
            delegate_strategy: The strategy to delegate actual search to
            all_links_of_system: Accumulated links from past searches
            settings_snapshot: Settings configuration
            research_context: Context from previous research
        """
        super().__init__(all_links_of_system, settings_snapshot)

        self.model = model
        self.search = search
        self.delegate_strategy = delegate_strategy
        self.research_context = research_context or {}

        # Initialize components
        self.relevance_filter = FollowUpRelevanceFilter(model)
        self.context_manager = FollowUpContextHandler(model)

        # Initialize question generator for creating contextualized queries
        self.question_generator = SimpleFollowUpQuestionGenerator(model)

        # For follow-up research, we ALWAYS want to combine sources
        # This is the whole point of follow-up - building on previous research
        self.combine_sources = True

        # Build comprehensive context
        self.full_context = self._build_full_context()

        logger.info(
            f"EnhancedContextualFollowUpStrategy initialized with "
            f"{len(self.full_context.get('past_sources', []))} past sources"
        )

    def _build_full_context(self) -> Dict[str, Any]:
        """
        Build comprehensive context from research data.

        Returns:
            Full context dictionary
        """
        # Use context manager to build structured context
        if self.research_context:
            # The follow-up query will be passed to analyze_topic later
            # For now, use empty string since we're just building initial context
            follow_up_query = ""
            context = self.context_manager.build_context(
                self.research_context, follow_up_query
            )

            # Also ensure we get the sources from the research context
            if "past_sources" not in context or not context["past_sources"]:
                # Try to get from various fields in research_context
                sources = []
                for field in ["resources", "all_links_of_system", "past_links"]:
                    if (
                        field in self.research_context
                        and self.research_context[field]
                    ):
                        sources.extend(self.research_context[field])
                context["past_sources"] = sources

            # Ensure we have the findings
            if "past_findings" not in context or not context["past_findings"]:
                context["past_findings"] = self.research_context.get(
                    "past_findings", ""
                )

            # Ensure we have the original query
            context["original_query"] = self.research_context.get(
                "original_query", ""
            )
        else:
            context = {
                "past_sources": [],
                "past_findings": "",
                "summary": "",
                "key_entities": [],
                "all_links_of_system": [],
                "original_query": "",
            }

        return context

    def analyze_topic(self, query: str) -> Dict:
        """
        Analyze a follow-up topic with enhanced context processing.

        This strategy:
        1. Reformulates the question based on past findings
        2. Filters relevant past sources to reuse
        3. Hands over to the delegate strategy with enhanced context

        Args:
            query: The follow-up question to research

        Returns:
            Research findings with context enhancement
        """
        logger.info(f"Starting enhanced follow-up search for: {query}")

        # Update the context with the actual follow-up query
        self.full_context["follow_up_query"] = query

        # Log what context we have
        logger.info(
            f"Context summary: {len(self.full_context.get('past_sources', []))} past sources, "
            f"findings length: {len(self.full_context.get('past_findings', ''))}, "
            f"original query: {self.full_context.get('original_query', 'N/A')}"
        )

        self._update_progress(
            "Analyzing past research context", 10, {"phase": "context_analysis"}
        )

        # Step 1: Skip reformulation - we'll use the original query with full context
        # This avoids LLM misinterpretation of queries like "provide data in a table"
        # The context will make it clear what the query refers to

        self._update_progress(
            "Preparing contextualized query",
            20,
            {"phase": "context_preparation", "original_query": query},
        )

        # Step 2: Filter relevant sources from past research using original query
        relevant_sources = self._filter_relevant_sources(query)

        self._update_progress(
            f"Identified {len(relevant_sources)} relevant past sources",
            30,
            {
                "phase": "source_filtering",
                "relevant_sources": len(relevant_sources),
                "total_past_sources": len(
                    self.full_context.get("past_sources", [])
                ),
            },
        )

        # Step 3: Inject the relevant sources into delegate strategy
        # This gives the delegate strategy a head start with pre-filtered sources
        self._inject_context_into_delegate(relevant_sources, query)

        self._update_progress(
            "Handing over to research strategy",
            40,
            {"phase": "delegate_handover"},
        )

        # Step 4: Create a query that includes FULL context from previous research
        # Use question generator to create contextualized query
        past_findings = self.full_context.get("past_findings", "")
        original_research_query = self.full_context.get("original_query", "")

        contextualized_query = (
            self.question_generator.generate_contextualized_query(
                follow_up_query=query,
                original_query=original_research_query,
                past_findings=past_findings,
            )
        )

        # Let the delegate strategy (from user's settings) do the actual research
        # with the contextualized query and pre-injected sources
        result = self.delegate_strategy.analyze_topic(contextualized_query)

        # Step 5: Get past sources for metadata (always needed)
        all_past_sources = self.full_context.get("past_sources", [])

        # Step 6: Optionally combine old sources with new ones
        # Only do this if the setting is enabled to avoid breaking existing reports
        if self.combine_sources:
            logger.info(
                f"Combining sources: self.combine_sources={self.combine_sources}"
            )

            # Ensure we have all sources from both researches
            if "all_links_of_system" not in result:
                result["all_links_of_system"] = []

            # Log initial state
            new_sources_count = len(result.get("all_links_of_system", []))
            logger.info(
                f"Initial state: {new_sources_count} new sources, {len(all_past_sources)} past sources to combine"
            )

            # Create a set of URLs already in the result to avoid duplicates
            existing_urls = {
                link.get("url")
                for link in result.get("all_links_of_system", [])
            }

            # Add all past sources that aren't already in the result
            added_count = 0
            for source in all_past_sources:
                url = source.get("url")
                if url and url not in existing_urls:
                    # Mark it as from previous research
                    enhanced_source = source.copy()
                    enhanced_source["from_past_research"] = True
                    result["all_links_of_system"].append(enhanced_source)
                    existing_urls.add(url)
                    added_count += 1

            logger.info(
                f"Source combination complete: Added {added_count} past sources, total now {len(result['all_links_of_system'])} sources"
            )
        else:
            logger.info(
                f"Source combination skipped: self.combine_sources={self.combine_sources}"
            )

        # Step 7: Add metadata about the follow-up enhancement
        if "metadata" not in result:
            result["metadata"] = {}

        result["metadata"]["follow_up_enhancement"] = {
            "original_query": query,
            "contextualized": True,
            "sources_reused": len(relevant_sources),
            "total_past_sources": len(all_past_sources),
            "parent_research_id": self.full_context.get(
                "parent_research_id", ""
            ),
        }

        self._update_progress(
            "Enhanced follow-up search complete",
            100,
            {
                "phase": "complete",
                "sources_reused": len(relevant_sources),
                "total_sources": len(result.get("all_links_of_system", [])),
            },
        )

        logger.info(
            f"Enhanced results: {len(relevant_sources)} past sources reused, "
            f"{len(result.get('all_links_of_system', []))} total sources found"
        )

        return result

    def _filter_relevant_sources(self, query: str) -> List[Dict]:
        """
        Filter past sources for relevance to the follow-up query.

        Args:
            query: The reformulated follow-up query

        Returns:
            List of relevant sources
        """
        past_sources = self.full_context.get("past_sources", [])

        if not past_sources:
            return []

        # Filter sources using the relevance filter
        # Get max sources from settings or use default
        max_followup_sources = self.settings_snapshot.get(
            "search.max_followup_sources", {}
        ).get("value", 15)

        relevant = self.relevance_filter.filter_results(
            results=past_sources,
            query=query,
            max_results=max_followup_sources,
            threshold=0.3,
            past_findings=self.full_context.get("past_findings", ""),
            original_query=self.full_context.get("original_query", ""),
        )

        logger.info(
            f"Filtered {len(past_sources)} past sources to "
            f"{len(relevant)} relevant ones"
        )

        return relevant

    def _inject_context_into_delegate(
        self, relevant_sources: List[Dict], reformulated_query: str
    ):
        """
        Inject context and sources into the delegate strategy.

        Args:
            relevant_sources: Filtered relevant sources
            reformulated_query: The reformulated query
        """
        # Initialize delegate's all_links_of_system if needed
        if not self.delegate_strategy.all_links_of_system:
            self.delegate_strategy.all_links_of_system = []

        # Add relevant sources to the beginning (high priority)
        existing_urls = {
            link.get("url")
            for link in self.delegate_strategy.all_links_of_system
        }

        injected_count = 0
        for source in relevant_sources:
            url = source.get("url")
            if url and url not in existing_urls:
                # Add source with enhanced metadata
                enhanced_source = source.copy()
                enhanced_source["from_past_research"] = True
                enhanced_source["follow_up_relevance"] = source.get(
                    "relevance_score", 1.0
                )

                self.delegate_strategy.all_links_of_system.insert(
                    0, enhanced_source
                )
                existing_urls.add(url)
                injected_count += 1

        logger.info(
            f"Injected {injected_count} relevant past sources into delegate strategy"
        )

        # Pass context to delegate if it supports it
        if hasattr(self.delegate_strategy, "set_followup_context"):
            self.delegate_strategy.set_followup_context(
                {
                    "reformulated_query": reformulated_query,
                    "past_findings_summary": self.full_context.get(
                        "summary", ""
                    ),
                    "key_entities": self.full_context.get("key_entities", []),
                    "sources_injected": injected_count,
                }
            )

    def set_progress_callback(self, callback):
        """
        Set progress callback for both wrapper and delegate.

        Args:
            callback: Progress callback function
        """
        super().set_progress_callback(callback)
        if self.delegate_strategy:
            self.delegate_strategy.set_progress_callback(callback)

    @property
    def citation_handler(self):
        """Expose the delegate's citation_handler.

        External code (notably ``research_service.py`` chat-mode streaming
        hookup) wires the stream callback via
        ``strategy.citation_handler.set_stream_callback(...)``. The wrapper
        itself owns no citation work — it delegates entirely to
        ``self.delegate_strategy`` — so transparently forwarding the
        attribute restores the streaming path that would otherwise be
        silently skipped for follow-up research.

        Returns ``None`` if no delegate is set (partial construction) so
        the ``hasattr`` check at the call site short-circuits cleanly.
        """
        if self.delegate_strategy is None:
            return None
        return getattr(self.delegate_strategy, "citation_handler", None)

    def close(self):
        """Forward close() to the delegate so resource-holding strategies
        (e.g. ConstraintParallelStrategy with its ThreadPoolExecutors) get
        a chance to shut down. Without this, AdvancedSearchSystem.close()
        would call BaseSearchStrategy.close() on the wrapper — a no-op —
        and never reach the delegate."""
        super().close()
        if self.delegate_strategy is not None:
            from ....utilities.resource_utils import safe_close

            safe_close(self.delegate_strategy, "delegate strategy")
