"""
IterDRAG strategy implementation.
"""

import json
from datetime import datetime, UTC
from typing import Dict, List

from loguru import logger

from ...citation_handler import CitationHandler

# Model and search should be provided by AdvancedSearchSystem
from ...config.thread_settings import get_setting_from_snapshot
from ...utilities.search_utilities import extract_links_from_search_results
from ..findings.repository import FindingsRepository
from ..knowledge.standard_knowledge import StandardKnowledge
from ..questions.decomposition_question import DecompositionQuestionGenerator
from .base_strategy import BaseSearchStrategy


class IterDRAGStrategy(BaseSearchStrategy):
    """IterDRAG strategy that breaks queries into sub-queries."""

    def __init__(
        self,
        search,
        model,
        max_iterations=3,
        subqueries_per_iteration=2,
        all_links_of_system=None,
        settings_snapshot=None,
    ):
        """Initialize the IterDRAG strategy with search and LLM.

        Args:
            search: Search engine to use for web queries
            model: LLM to use for text generation and reasoning
            max_iterations: Maximum number of iterations to run
            subqueries_per_iteration: Number of sub-queries to generate per iteration
            all_links_of_system: Optional list of links to initialize with
            settings_snapshot: Settings snapshot for thread context
        """
        super().__init__(
            all_links_of_system=all_links_of_system,
            settings_snapshot=settings_snapshot,
        )
        self.search = search
        self.model = model
        self.max_iterations = max_iterations
        self.subqueries_per_iteration = subqueries_per_iteration

        # Initialize progress callback
        self.progress_callback = None
        # Note: questions_by_iteration is already initialized by parent class

        # Create citation handler with settings
        self.citation_handler = CitationHandler(
            self.model, settings_snapshot=settings_snapshot
        )

        # Initialize components
        self.question_generator = DecompositionQuestionGenerator(self.model)
        self.knowledge_generator = StandardKnowledge(self.model)
        self.findings_repository = FindingsRepository(self.model)

    def _generate_subqueries(
        self, query: str, initial_results: List[Dict], current_knowledge: str
    ) -> List[str]:
        """Generate sub-queries based on initial search results and the main query."""
        try:
            # Format context for question generation
            context = f"""Current Query: {query}
Current Date: {datetime.now(UTC).strftime("%Y-%m-%d")}
Past Questions: {self.questions_by_iteration}
Current Knowledge: {current_knowledge}

Initial Search Results:
{json.dumps(initial_results, indent=2)}"""

            # Generate sub-queries using the question generator
            return self.question_generator.generate_questions(
                query,
                context,
                questions_per_iteration=int(
                    get_setting_from_snapshot(
                        "search.questions_per_iteration",
                        settings_snapshot=self.settings_snapshot,
                    )
                    or 3
                ),
            )
        except Exception:
            logger.exception("Error generating sub-queries")
            return []

    def analyze_topic(self, query: str) -> Dict:
        """IterDRAG implementation of the topic analysis process."""
        findings = []
        current_knowledge = ""
        section_links = []

        self._update_progress(
            "Initializing IterDRAG research system",
            5,
            {"phase": "init", "strategy": "iterdrag"},
        )

        # Check if search engine is available
        if self.search is None:
            error_msg = "Error: No search engine available. Please check your configuration."
            self._update_progress(
                error_msg,
                100,
                {
                    "phase": "error",
                    "error": "No search engine available",
                    "status": "failed",
                },
            )
            return {
                "findings": [],
                "iterations": 0,
                "questions": {},
                "formatted_findings": "Error: Unable to conduct research without a search engine.",
                "current_knowledge": "",
                "error": error_msg,
            }

        # Initial search for the main query
        self._update_progress(
            "Performing initial search for main query",
            10,
            {"phase": "search", "iteration": 1},
        )

        initial_results = self.search.run(query)
        if not initial_results:
            self._update_progress(
                "No initial results found",
                15,
                {"phase": "search_complete", "result_count": 0},
            )
            initial_results = []
        else:
            self._update_progress(
                f"Found {len(initial_results)} initial results",
                15,
                {
                    "phase": "search_complete",
                    "result_count": len(initial_results),
                },
            )

        # Extract and save links
        initial_links = extract_links_from_search_results(initial_results)
        self.all_links_of_system.extend(initial_links)
        section_links.extend(initial_links)

        # Generate sub-queries
        self._update_progress(
            "Generating sub-queries for IterDRAG analysis",
            20,
            {"phase": "iterdrag_decomposition"},
        )

        sub_queries = self._generate_subqueries(
            query, initial_results, current_knowledge
        )

        # Store questions in repository and in self.questions_by_iteration
        self.questions_by_iteration = {0: sub_queries}
        self.findings_repository.set_questions_by_iteration(
            self.questions_by_iteration
        )

        if not sub_queries:
            logger.error("No sub-queries were generated to analyze.")
            finding = {
                "phase": "Analysis Error",
                "content": "No sub-queries could be generated for the main question.",
                "question": query,
                "search_results": [],
            }
            findings.append(finding)
        else:
            # Process each sub-query
            total_subqueries = len(sub_queries)

            for i, sub_query in enumerate(sub_queries, 1):
                progress_base = 25 + (i / total_subqueries * 50)
                self._update_progress(
                    f"Processing sub-query {i} of {total_subqueries}: {sub_query}",
                    int(progress_base),
                    {"phase": "subquery", "subquery_index": i},
                )

                # Search for the sub-query
                try:
                    sub_results = self.search.run(sub_query)
                    if not sub_results:
                        self._update_progress(
                            f"No results for sub-query: {sub_query}",
                            int(progress_base + 2),
                            {"phase": "search_complete", "result_count": 0},
                        )
                        sub_results = []
                    else:
                        self._update_progress(
                            f"Found {len(sub_results)} results for sub-query: {sub_query}",
                            int(progress_base + 2),
                            {
                                "phase": "search_complete",
                                "result_count": len(sub_results),
                            },
                        )
                except Exception:
                    logger.exception("Error searching for sub-query")
                    sub_results = []

                try:
                    # Use previous knowledge to answer this sub-query
                    result = self.citation_handler.analyze_followup(
                        sub_query,
                        sub_results,
                        current_knowledge,
                        nr_of_links=len(self.all_links_of_system),
                    )

                    # Extract and save links AFTER citation handler processes results
                    sub_links = extract_links_from_search_results(sub_results)
                    self.all_links_of_system.extend(sub_links)
                    section_links.extend(sub_links)

                    if result is not None:
                        finding = {
                            "phase": f"Sub-query {i}",
                            "content": result["content"],
                            "question": sub_query,
                            "search_results": sub_results,
                            "documents": result["documents"],
                        }
                        findings.append(finding)
                        self.findings_repository.add_finding(sub_query, finding)
                        self.findings_repository.add_documents(
                            result["documents"]
                        )

                        # Add to current knowledge with space around +
                        current_knowledge = (
                            current_knowledge
                            + "\n\n\n New: \n"
                            + result["content"]
                        )
                except Exception:
                    logger.exception("Error analyzing sub-query results:")
                    finding = {
                        "phase": f"Follow-up Iteration 0.{i + 1}",
                        "content": "Error analyzing sub-query results.",
                        "question": sub_query,
                        "search_results": [],
                        "documents": [],
                    }
                    findings.append(finding)
                # Optionally update current_knowledge even if analysis failed?

            # Final answer synthesis based on all sub-query findings
            self._update_progress(
                "Synthesizing final answer from sub-queries",
                80,
                {"phase": "final_synthesis"},
            )

            try:
                # Extract finding contents for synthesis
                finding_contents = [
                    f["content"] for f in findings if "content" in f
                ]

                # Synthesize findings into a final answer
                final_answer = self.findings_repository.synthesize_findings(
                    query=query,
                    sub_queries=sub_queries,
                    findings=finding_contents,
                    accumulated_knowledge=current_knowledge,
                )

                # Check if the synthesis failed with an error
                if isinstance(final_answer, str) and final_answer.startswith(
                    "Error:"
                ):
                    logger.error(f"Synthesis returned an error: {final_answer}")

                    # Extract error type for better handling
                    error_type = "unknown"
                    error_message = final_answer.lower()

                    if "timeout" in error_message:
                        error_type = "timeout"
                    elif "token limit" in error_message:
                        error_type = "token_limit"
                    elif "rate limit" in error_message:
                        error_type = "rate_limit"
                    elif "connection" in error_message:
                        error_type = "connection"

                    # Log with specific error type
                    logger.error(
                        f"Synthesis failed with error type: {error_type}"
                    )

                    # Add error information to progress update
                    self._update_progress(
                        f"Synthesis failed: {error_type}. Attempting fallback...",
                        85,
                        {"phase": "synthesis_error", "error_type": error_type},
                    )

                    # Try fallback: use the best individual finding as the final answer
                    longest_finding = ""
                    for f in findings:
                        if isinstance(f, dict) and "content" in f:
                            if len(f["content"]) > len(longest_finding):
                                longest_finding = f["content"]

                    if longest_finding:
                        logger.info(
                            "Using longest finding as fallback synthesis"
                        )
                        final_answer = f"""
# Research Results (Fallback Mode)

{longest_finding}

## Note
The final synthesis could not be completed due to an error: {final_answer}
This is a fallback response using the most detailed individual finding.
                        """
                    else:
                        # If we don't have any findings with content, use current_knowledge
                        logger.info(
                            "Using current knowledge as fallback synthesis"
                        )
                        final_answer = f"""
# Research Results (Fallback Mode)

{current_knowledge}

## Note
The final synthesis could not be completed due to an error: {final_answer}
This is a fallback response using the accumulated knowledge.
                        """

                # Create a synthesis finding
                finding = {
                    "phase": "Final synthesis",
                    "content": final_answer,
                    "question": query,
                    "search_results": [],
                    "documents": [],
                }
                findings.append(finding)

                # Store the synthesized content
                self.findings_repository.add_finding(
                    query + "_synthesis", final_answer
                )

                # Update current knowledge with the synthesized version
                current_knowledge = final_answer
            except Exception as e:
                logger.exception("Error synthesizing final answer")

                # Create an error finding
                error_finding = {
                    "phase": "Final synthesis error",
                    "content": f"Error during synthesis: {e!s}",
                    "question": query,
                    "search_results": [],
                    "documents": [],
                }
                findings.append(error_finding)

                # If synthesis completely fails, construct a fallback answer from the most relevant findings
                self._update_progress(
                    "Synthesis failed. Creating fallback summary...",
                    85,
                    {"phase": "synthesis_fallback"},
                )

                try:
                    # Extract best content from findings
                    key_findings = []
                    for i, f in enumerate(findings):
                        if (
                            isinstance(f, dict)
                            and "content" in f
                            and f.get("content")
                        ):
                            # Only take the first 500 chars of each finding for the fallback
                            content_preview = f.get("content", "")[:500]
                            if content_preview:
                                key_findings.append(
                                    f"### Finding {i + 1}\n\n{content_preview}..."
                                )

                    # Create fallback content
                    fallback_content = f"""
# Research Results (Error Recovery Mode)

## Original Query
{query}

## Key Findings
{chr(10).join(key_findings[:5]) if key_findings else "No valid findings were generated."}

## Error Information
The system encountered an error during final synthesis: {e!s}
This is an automatically generated fallback response.
                    """

                    final_answer = fallback_content
                except Exception as fallback_error:
                    # Last resort fallback
                    logger.exception("Even fallback creation failed")
                    final_answer = f"""
# Research Error

The system encountered multiple errors while processing your query: "{query}"

Primary error: {e!s}
Fallback error: {fallback_error!s}

Please try again with a different query or contact support.
                    """

        # Compress knowledge if needed
        if (
            get_setting_from_snapshot(
                "general.knowledge_accumulation",
                "ITERATION",
                settings_snapshot=self.settings_snapshot,
            )
            == "ITERATION"
        ):
            try:
                self._update_progress(
                    "Compressing knowledge",
                    90,
                    {"phase": "knowledge_compression"},
                )
                current_knowledge = self.knowledge_generator.compress_knowledge(
                    current_knowledge, query, section_links
                )
            except Exception:
                logger.exception("Error compressing knowledge")

        # Format and save findings
        self._update_progress(
            "Formatting and saving findings", 95, {"phase": "formatting"}
        )

        try:
            # First, get just the synthesized content without formatting
            if not isinstance(final_answer, str):
                logger.error(
                    "final_answer is not a string, using current_knowledge as fallback"
                )
                final_answer = current_knowledge

            # Ensure latest questions are in the repository
            self.findings_repository.set_questions_by_iteration(
                self.questions_by_iteration
            )

            # Now format the findings with search questions and sources
            formatted_findings = (
                self.findings_repository.format_findings_to_text(
                    findings, final_answer
                )
            )
        except Exception:
            logger.exception("Error formatting final findings")
            formatted_findings = (
                "Error: Could not format findings due to an error."
            )

        self._update_progress("Research complete", 100, {"phase": "complete"})

        return {
            "findings": findings,  # Keep the detailed findings list
            "iterations": 1,
            "questions": self.questions_by_iteration,  # Use the member variable instead of {"0": sub_queries}
            "formatted_findings": formatted_findings,  # This is the fully formatted string for UI
            "current_knowledge": current_knowledge,  # This is the synthesized content for knowledge base
        }
