"""
Iterative Refinement Strategy

This strategy orchestrates iterative refinement of research results by:
1. Running initial research with a base strategy
2. Using LLM to evaluate results and identify gaps
3. Generating follow-up queries to address gaps
4. Calling the follow-up strategy to refine results
5. Repeating until quality threshold is met or max iterations reached
"""

from typing import Dict, Optional
from loguru import logger

from ...utilities.json_utils import extract_json, get_llm_response_text

from .base_strategy import BaseSearchStrategy
from .followup.enhanced_contextual_followup import (
    EnhancedContextualFollowUpStrategy,
)


class IterativeRefinementStrategy(BaseSearchStrategy):
    """
    Strategy that iteratively refines research results using LLM-guided decisions.

    This strategy acts as an orchestrator that:
    - Runs initial research using any base strategy
    - Evaluates results quality with LLM
    - Generates targeted follow-up queries
    - Uses EnhancedContextualFollowUpStrategy for refinement
    - Accumulates knowledge across iterations
    """

    def __init__(
        self,
        model,
        search,
        initial_strategy: BaseSearchStrategy,
        all_links_of_system=None,
        settings_snapshot=None,
        evaluation_frequency: int = 1,
        max_refinements: int = 3,
        confidence_threshold: float = 0.8,
        max_evaluation_tokens: Optional[int] = None,  # None = no truncation
        **kwargs,
    ):
        """
        Initialize the iterative refinement strategy.

        Args:
            model: The LLM model to use
            search: The search engine
            initial_strategy: Base strategy for initial research and follow-up delegate
            all_links_of_system: Accumulated links from past searches
            settings_snapshot: Settings configuration
            evaluation_frequency: Evaluate after N refinement rounds (1 = every round)
            max_refinements: Maximum number of refinement cycles
            confidence_threshold: Confidence level to stop refinement (0.0-1.0)
            max_evaluation_tokens: Maximum tokens for LLM evaluation (None = no truncation)
        """
        super().__init__(all_links_of_system, settings_snapshot)

        # Validate parameters
        if not 0 <= confidence_threshold <= 1:
            raise ValueError("confidence_threshold must be between 0 and 1")
        if max_refinements < 0:
            raise ValueError("max_refinements must be non-negative")
        if evaluation_frequency < 1:
            raise ValueError("evaluation_frequency must be at least 1")
        if max_evaluation_tokens is not None and max_evaluation_tokens <= 0:
            raise ValueError("max_evaluation_tokens must be positive or None")

        self.model = model
        self.search = search
        self.initial_strategy = initial_strategy
        self.evaluation_frequency = evaluation_frequency
        self.max_refinements = max_refinements
        self.confidence_threshold = confidence_threshold
        self.max_evaluation_tokens = max_evaluation_tokens

        # Track refinement history
        self.refinement_history = []

        logger.info(
            f"IterativeRefinementStrategy initialized with "
            f"max_refinements={max_refinements}, "
            f"confidence_threshold={confidence_threshold}, "
            f"evaluation_frequency={evaluation_frequency}"
        )

    def analyze_topic(self, query: str) -> Dict:
        """
        Analyze a topic with iterative refinement.

        Args:
            query: The research query to analyze

        Returns:
            Research findings with iterative refinements
        """
        logger.info(f"Starting iterative refinement research for: {query}")

        self._update_progress(
            "Running initial research",
            10,
            {"phase": "initial_research", "query": query},
        )

        # Step 1: Run initial research
        results = self.initial_strategy.analyze_topic(query)

        # Validate initial results before proceeding
        if not results:
            logger.warning("Initial research returned no results")
            return {
                "findings": [],
                "iterations": 0,
                "questions_by_iteration": {},
                "formatted_findings": "Error: Initial research failed to return results.",
                "current_knowledge": "",
                "error": "Initial research failed",
            }

        if "error" in results:
            logger.error(
                f"Initial research failed with error: {results['error']}"
            )
            return results

        if not results.get("all_links_of_system"):
            logger.warning(
                "Initial research returned no sources, skipping refinement"
            )
            # Still return the results we have, just without refinement
            results["refinement_metadata"] = {
                "strategy": "iterative_refinement",
                "refinements_performed": 0,
                "max_refinements": self.max_refinements,
                "final_confidence": 0.0,
                "total_sources": 0,
                "refinement_history": [],
                "skipped_reason": "No initial sources found",
            }
            return results

        # Build initial research context
        research_context = self._build_research_context(results, query)

        # Track all accumulated sources
        all_accumulated_sources = results.get("all_links_of_system", []).copy()

        refinement_count = 0
        rounds_since_evaluation = 0

        # Main refinement loop
        while refinement_count < self.max_refinements:
            self.check_termination()
            rounds_since_evaluation += 1

            # Evaluate if it's time (based on frequency)
            if rounds_since_evaluation >= self.evaluation_frequency:
                progress = 20 + (refinement_count * 60 / self.max_refinements)

                self._update_progress(
                    f"Evaluating research quality (refinement {refinement_count + 1}/{self.max_refinements})",
                    progress,
                    {"phase": "evaluation", "refinement": refinement_count + 1},
                )

                # Step 2: Use LLM to evaluate current results
                evaluation = self._evaluate_with_llm(
                    results, query, refinement_count
                )

                logger.info(
                    f"LLM evaluation - Action: {evaluation.get('action')}, "
                    f"Confidence: {evaluation.get('confidence', 0):.2f}"
                )

                # Record evaluation in history with full details
                self.refinement_history.append(
                    {
                        "iteration": refinement_count + 1,
                        "evaluation": evaluation,
                        "sources_count": len(all_accumulated_sources),
                        "action": evaluation.get("action"),
                        "confidence": evaluation.get("confidence", 0),
                        "gaps": evaluation.get("gaps", []),
                        "reasoning": evaluation.get("reasoning", ""),
                        "query": evaluation.get("refinement_query", ""),
                    }
                )

                # Check if we should stop
                if evaluation.get("action") == "COMPLETE":
                    logger.info("LLM determined research is complete")
                    break

                if evaluation.get("confidence", 0) >= self.confidence_threshold:
                    logger.info(
                        f"Confidence threshold met: {evaluation.get('confidence'):.2f} >= {self.confidence_threshold}"
                    )
                    break

                # Step 3: Generate follow-up query for refinement
                followup_query = evaluation.get("refinement_query")

                if not followup_query:
                    logger.warning(
                        "No refinement query generated, stopping refinement"
                    )
                    break

                logger.info(f"Refinement query: {followup_query}")

                self._update_progress(
                    f"Running follow-up research: {followup_query[:100]}...",
                    progress + 10,
                    {"phase": "followup", "query": followup_query},
                )

                # Step 4: Create follow-up strategy with accumulated context
                # Reuse the initial strategy directly - it already has all configuration
                # The EnhancedContextualFollowUpStrategy will handle context properly
                followup_strategy = EnhancedContextualFollowUpStrategy(
                    model=self.model,
                    search=self.search,
                    delegate_strategy=self.initial_strategy,
                    all_links_of_system=[],  # Fresh for this iteration
                    settings_snapshot=self.settings_snapshot,
                    research_context=research_context,
                )

                # Pass progress callback if available
                if self.progress_callback:
                    followup_strategy.set_progress_callback(
                        self.progress_callback
                    )

                # Step 5: Run follow-up research
                followup_results = followup_strategy.analyze_topic(
                    followup_query
                )

                # Check if follow-up produced useful results
                new_sources = followup_results.get("all_links_of_system", [])
                if len(new_sources) < 2:
                    logger.info(
                        f"Refinement yielded only {len(new_sources)} new sources, stopping refinement"
                    )
                    break

                # Step 6: Merge results
                results = self._merge_results(
                    results, followup_results, followup_query
                )

                # Update accumulated sources (O(1) lookup with set)
                existing_urls = {s.get("url") for s in all_accumulated_sources}
                new_sources = followup_results.get("all_links_of_system", [])
                for source in new_sources:
                    # Check if URL already exists to avoid duplicates
                    if source.get("url") not in existing_urls:
                        all_accumulated_sources.append(source)
                        existing_urls.add(source.get("url"))

                # Update research context for next iteration
                research_context = self._build_research_context(results, query)

                # Reset evaluation counter
                rounds_since_evaluation = 0

            refinement_count += 1

        # Reorganize formatted_findings to show final result at the top
        separator_check = "---\n## Refinement:" in results.get(
            "formatted_findings", ""
        )
        logger.info(
            f"Checking reorganization: refinement_count={refinement_count}, has separator={separator_check}"
        )
        if refinement_count > 0 and "---\n## Refinement:" in results.get(
            "formatted_findings", ""
        ):
            logger.info(
                f"Reorganizing findings with {refinement_count} refinements"
            )
            # Split to get the last refinement
            parts = results["formatted_findings"].rsplit(
                "---\n## Refinement:", 1
            )
            if len(parts) == 2:
                last_refinement = "## Refinement:" + parts[1]
                # Extract just the content after the query line
                if "\n" in last_refinement:
                    _, final_content = last_refinement.split("\n", 1)

                    # Build confidence progression and evaluation details
                    confidence_progression = []
                    for hist in self.refinement_history:
                        conf = hist["evaluation"].get("confidence", 0)
                        confidence_progression.append(f"{conf:.0%}")

                    confidence_str = (
                        " → ".join(confidence_progression)
                        if confidence_progression
                        else "N/A"
                    )

                    # Build evaluation insights section
                    evaluation_insights = []
                    for i, hist in enumerate(self.refinement_history, 1):
                        gaps = hist.get("gaps", [])
                        reasoning = hist.get("reasoning", "")
                        query = hist.get("query", "")
                        action = hist.get("action", "")

                        if gaps or reasoning:
                            evaluation_insights.append(f"""
**Refinement {i} Analysis:**
- **Action**: {action}
- **Gaps Identified**: {", ".join(gaps) if gaps else "None"}
- **Reasoning**: {reasoning}
- **Follow-up Query**: {query if query else "N/A"}""")

                    insights_section = ""
                    if evaluation_insights:
                        insights_section = f"""

---

## Refinement Analysis

The LLM identified these issues and improvements during the refinement process:
{"".join(evaluation_insights)}"""

                    # Prepend final answer at top
                    results["formatted_findings"] = f"""## Final Answer
*After {refinement_count} refinement{"s" if refinement_count != 1 else ""} | Confidence progression: {confidence_str}*

{final_content}{insights_section}

---

## Research Progression

{results["formatted_findings"]}"""

        # Add refinement metadata to results
        results["refinement_metadata"] = {
            "strategy": "iterative_refinement",
            "refinements_performed": refinement_count,
            "max_refinements": self.max_refinements,
            "final_confidence": self.refinement_history[-1]["evaluation"].get(
                "confidence", 0
            )
            if self.refinement_history
            else 0,
            "total_sources": len(all_accumulated_sources),
            "refinement_history": self.refinement_history,
        }

        # Ensure all accumulated sources are in final results
        results["all_links_of_system"] = all_accumulated_sources

        self._update_progress(
            f"Research complete after {refinement_count} refinement(s)",
            100,
            {"phase": "complete", "refinements": refinement_count},
        )

        logger.info(
            f"Iterative refinement complete: {refinement_count} refinements, "
            f"{len(all_accumulated_sources)} total sources"
        )

        return results

    def _evaluate_with_llm(
        self, results: Dict, original_query: str, iteration: int
    ) -> Dict:
        """
        Use LLM to evaluate current research results and decide on refinement.

        Args:
            results: Current research results
            original_query: The original research query
            iteration: Current refinement iteration number

        Returns:
            Dictionary with evaluation results including action and refinement query
        """
        # Get formatted findings for evaluation
        findings = results.get("formatted_findings", "")
        sources_count = len(results.get("all_links_of_system", []))

        # Build gaps summary from refinement history
        previous_gaps = []
        if self.refinement_history:
            for hist in self.refinement_history:
                gaps = hist["evaluation"].get("gaps", [])
                previous_gaps.extend(gaps)

        prompt = f"""You are evaluating research results to determine if refinement is needed.

Original Research Query: {original_query}
Current Refinement Iteration: {iteration + 1} of {self.max_refinements}

Current Research Findings:
{findings if self.max_evaluation_tokens is None else findings[: self.max_evaluation_tokens]}

Number of sources found: {sources_count}
Previous gaps identified: {", ".join(previous_gaps) if previous_gaps else "None"}

Your task is to critically evaluate these research results and decide what would most improve them.

Key questions to consider:
- Are the main entities/companies mentioned still active and correctly described?
- Is the information current and up-to-date?
- Are we missing obvious major players or important aspects?
- Are there any suspicious claims that need verification?
- Does the answer fully address what was asked?
- Are there contradictions or inconsistencies?

You have complete freedom to:
- Ask verification questions ("Is X company still active in this field?")
- Challenge assumptions ("Shouldn't Y also be considered a leader?")
- Request missing perspectives ("What about competitor Z's approach?")
- Seek recent updates ("What's the latest status as of 2024/2025?")
- Explore contradictions ("Source says X but that conflicts with Y")

Respond with ONLY valid JSON, no markdown formatting, no additional text before or after:
{{
    "action": "REFINE" or "COMPLETE",
    "confidence": 0.0 to 1.0 (your honest assessment of completeness/accuracy),
    "gaps": ["list", "of", "identified", "gaps", "or", "concerns"],
    "reasoning": "Your honest evaluation - what's good, what's problematic, what's missing",
    "refinement_query": "If action is REFINE, ask whatever question would most improve the research - be creative and thorough"
}}

Be critical and thorough. It's better to verify and be sure than to present incorrect information.
Your refinement query can be as long and specific as needed to get good results."""

        try:
            response = self.model.invoke(prompt)
            content = get_llm_response_text(response)

            evaluation = extract_json(content, expected_type=dict)
        except Exception as e:
            logger.exception("Error in LLM evaluation")
            # Return default to continue without refinement
            return {
                "action": "COMPLETE",
                "confidence": 0.5,
                "gaps": [],
                "reasoning": f"Evaluation error: {str(e)}",
                "refinement_query": None,
            }

        if evaluation is None:
            return {
                "action": "COMPLETE",
                "confidence": 0.5,
                "gaps": [],
                "reasoning": "Could not parse JSON from LLM response",
                "refinement_query": None,
            }

        # Validate and clean up the evaluation
        if "action" not in evaluation:
            evaluation["action"] = "COMPLETE"

        evaluation["action"] = evaluation["action"].upper()
        if evaluation["action"] not in ["REFINE", "COMPLETE"]:
            evaluation["action"] = "COMPLETE"

        # Ensure confidence is a float between 0 and 1
        evaluation["confidence"] = max(
            0.0, min(1.0, float(evaluation.get("confidence", 0.5)))
        )

        # Ensure gaps is a list
        if not isinstance(evaluation.get("gaps", []), list):
            evaluation["gaps"] = []

        return evaluation

    def _build_research_context(
        self, results: Dict, original_query: str
    ) -> Dict:
        """
        Build research context for follow-up strategy.

        Args:
            results: Current research results
            original_query: Original research query

        Returns:
            Context dictionary for follow-up strategy
        """
        context = {
            "original_query": original_query,
            "past_findings": results.get("formatted_findings", ""),
            "past_sources": results.get("all_links_of_system", []),
            "resources": results.get("all_links_of_system", []),
            "delegate_strategy": self.initial_strategy.__class__.__name__.replace(
                "Strategy", ""
            )
            .replace("SearchStrategy", "")
            .lower(),
        }

        # Add iteration information if available
        if "questions_by_iteration" in results:
            context["questions_by_iteration"] = results[
                "questions_by_iteration"
            ]

        # Add findings if available
        if "findings" in results:
            context["findings"] = results["findings"]

        return context

    def _merge_results(
        self,
        original_results: Dict,
        followup_results: Dict,
        followup_query: str,
    ) -> Dict:
        """
        Merge follow-up results into original results.

        Args:
            original_results: Original research results
            followup_results: Follow-up research results
            followup_query: The follow-up query used

        Returns:
            Merged results dictionary
        """
        merged = original_results.copy()

        # Append new findings
        if "findings" in followup_results:
            if "findings" not in merged:
                merged["findings"] = []

            # Add a separator finding
            merged["findings"].append(
                {
                    "phase": f"Refinement: {followup_query[:100]}",
                    "content": f"Follow-up research for: {followup_query}",
                    "type": "refinement_separator",
                }
            )

            # Add the new findings
            merged["findings"].extend(followup_results["findings"])

        # Merge formatted findings with clear separation
        if "formatted_findings" in followup_results:
            original_findings = merged.get("formatted_findings", "")
            followup_findings = followup_results["formatted_findings"]

            merged["formatted_findings"] = f"""{original_findings}

---
## Refinement: {followup_query}

{followup_findings}"""

        # Update current knowledge to latest
        if "current_knowledge" in followup_results:
            merged["current_knowledge"] = followup_results["current_knowledge"]

        # Merge questions by iteration
        if "questions_by_iteration" in followup_results:
            if "questions_by_iteration" not in merged:
                merged["questions_by_iteration"] = {}

            # Add follow-up questions with new keys
            max_iter = (
                max(merged["questions_by_iteration"].keys())
                if merged["questions_by_iteration"]
                else 0
            )
            # Create a list copy of items to avoid dictionary changed size during iteration error
            followup_items = list(
                followup_results["questions_by_iteration"].items()
            )
            for key, questions in followup_items:
                merged["questions_by_iteration"][max_iter + key] = questions

        # Merge links (handled in main loop to avoid duplicates)
        # Don't merge here as we handle it in analyze_topic

        return merged

    def set_progress_callback(self, callback):
        """
        Set progress callback for both this strategy and delegates.

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