"""
Iterative Reasoning Strategy for step-by-step search and reasoning.

This strategy maintains a persistent knowledge base and iteratively:
1. Analyzes what we know so far
2. Decides what to search next
3. Performs the search
4. Updates knowledge with findings
5. Repeats until confident in answer
"""

from dataclasses import dataclass, field
from datetime import datetime, UTC
from typing import Any, Dict, List, Optional

from langchain_core.language_models import BaseChatModel
from loguru import logger

from ...utilities.search_utilities import format_findings, remove_think_tags
from ..findings.repository import FindingsRepository
from .base_strategy import BaseSearchStrategy
from .source_based_strategy import SourceBasedSearchStrategy


@dataclass
class KnowledgeState:
    """Current state of our knowledge about the query."""

    original_query: str
    key_facts: List[str] = field(default_factory=list)
    uncertainties: List[str] = field(default_factory=list)
    search_history: List[Dict] = field(default_factory=list)
    candidate_answers: List[Dict] = field(default_factory=list)
    confidence: float = 0.0
    iteration: int = 0

    def to_string(self) -> str:
        """Convert knowledge state to readable string for LLM."""
        return f"""
Original Query: {self.original_query}

What We Know:
{chr(10).join(f"- {fact}" for fact in self.key_facts) if self.key_facts else "- Nothing yet"}

What We're Uncertain About:
{chr(10).join(f"- {uncertainty}" for uncertainty in self.uncertainties) if self.uncertainties else "- Nothing specific"}

Search History ({len(self.search_history)} searches):
{chr(10).join(f"- Search {i + 1}: {s['query']}" for i, s in enumerate(self.search_history[-3:])) if self.search_history else "- No searches yet"}

Candidate Answers ({len(self.candidate_answers)} total):
{chr(10).join(f"- {c['answer']} (confidence: {c['confidence']:.0%})" for c in sorted(self.candidate_answers, key=lambda x: x["confidence"], reverse=True)) if self.candidate_answers else "- None yet"}

Current Confidence: {self.confidence:.1%}
"""


class IterativeReasoningStrategy(BaseSearchStrategy):
    """
    A strategy that iteratively searches and reasons, maintaining persistent knowledge.

    Simple loop:
    1. Assess current knowledge
    2. Decide next search
    3. Execute search
    4. Update knowledge
    5. Check if we have answer
    """

    def __init__(
        self,
        model: BaseChatModel,
        search: Any,
        all_links_of_system: List[str],
        max_iterations: int = 10,
        confidence_threshold: float = 0.85,
        search_iterations_per_round: int = 1,
        questions_per_search: int = 15,
    ):
        """Initialize the iterative reasoning strategy.

        Args:
            model: The language model to use
            search: The search engine instance
            all_links_of_system: List to store all encountered links
            max_iterations: Maximum reasoning iterations
            confidence_threshold: Confidence needed to stop
            search_iterations_per_round: Iterations per search round
            questions_per_search: Questions per search
        """
        super().__init__(all_links_of_system)
        self.model = model
        self.search = search
        self.max_iterations = max_iterations
        self.confidence_threshold = confidence_threshold
        self.search_iterations_per_round = search_iterations_per_round
        self.questions_per_search = questions_per_search
        self.findings_repository = FindingsRepository(model)
        self.knowledge_state: Optional[KnowledgeState] = None

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

        Args:
            query: The research query to analyze

        Returns:
            Dictionary containing analysis results
        """
        # Initialize
        self.all_links_of_system.clear()
        self.questions_by_iteration = []
        self.findings = []
        self.knowledge_state = KnowledgeState(original_query=query)

        # Progress callback
        if self.progress_callback:
            self.progress_callback(
                "Starting iterative reasoning - will build knowledge step by step",
                1,
                {
                    "phase": "init",
                    "strategy": "iterative_reasoning",
                    "reasoning_approach": "step-by-step",
                },
            )

        logger.info(f"Starting iterative reasoning for: {query}")
        logger.info(
            f"Max iterations: {self.max_iterations}, Confidence threshold: {self.confidence_threshold}"
        )

        # Add initial analysis to findings
        initial_finding = {
            "phase": "Initial Analysis",
            "content": f"""
**Query**: {query}

**Strategy**: Iterative Reasoning
- Will build knowledge step-by-step
- Continue until {self.confidence_threshold:.0%} confident or {self.max_iterations} steps
- Each step will search, analyze findings, and update knowledge

**Starting Analysis**...
""".strip(),
            "timestamp": self._get_timestamp(),
        }
        self.findings.append(initial_finding)

        # Main reasoning loop
        while (
            self.knowledge_state.iteration < self.max_iterations
            and self.knowledge_state.confidence < self.confidence_threshold
        ):
            self.knowledge_state.iteration += 1
            logger.info(
                f"Iteration {self.knowledge_state.iteration}/{self.max_iterations}, Current confidence: {self.knowledge_state.confidence:.1%}"
            )

            # Update progress
            if self.progress_callback:
                progress = (
                    int(
                        (self.knowledge_state.iteration / self.max_iterations)
                        * 80
                    )
                    + 10
                )
                facts_count = len(self.knowledge_state.key_facts)
                candidates_count = len(self.knowledge_state.candidate_answers)

                # Format top candidates for display
                top_candidates = ""
                if self.knowledge_state.candidate_answers:
                    sorted_candidates = sorted(
                        self.knowledge_state.candidate_answers,
                        key=lambda x: x["confidence"],
                        reverse=True,
                    )[:2]
                    top_candidates = ", ".join(
                        f"{c['answer']} ({c['confidence']:.0%})"
                        for c in sorted_candidates
                    )

                self.progress_callback(
                    f"Step {self.knowledge_state.iteration}: {facts_count} facts, {candidates_count} candidates. Top: {top_candidates if top_candidates else 'none yet'}",
                    progress,
                    {
                        "phase": "reasoning",
                        "iteration": self.knowledge_state.iteration,
                        "confidence": self.knowledge_state.confidence,
                        "facts_found": facts_count,
                        "uncertainties": len(
                            self.knowledge_state.uncertainties
                        ),
                        "candidates_count": candidates_count,
                        "top_candidates": top_candidates,
                    },
                )

            # Step 1: Analyze current knowledge and decide next search
            if self.progress_callback:
                self.progress_callback(
                    f"Step {self.knowledge_state.iteration}: Deciding what to search next",
                    progress + 2,
                    {"phase": "planning", "step": "deciding_search"},
                )

            next_search = self._decide_next_search()

            if not next_search:
                logger.info("No more searches needed")
                if self.progress_callback:
                    self.progress_callback(
                        "Sufficient information gathered - preparing final answer",
                        progress + 5,
                        {"phase": "synthesis", "reason": "no_more_searches"},
                    )
                break

            # Step 2: Execute the search
            if self.progress_callback:
                self.progress_callback(
                    f"Step {self.knowledge_state.iteration}: Searching: {next_search[:50]}...",
                    progress + 5,
                    {"phase": "searching", "query": next_search},
                )

            search_results = self._execute_search(next_search)

            # Step 3: Update knowledge with findings
            if self.progress_callback:
                self.progress_callback(
                    f"Step {self.knowledge_state.iteration}: Processing search results",
                    progress + 8,
                    {"phase": "processing", "step": "updating_knowledge"},
                )

            # Store confidence before update
            prev_confidence = self.knowledge_state.confidence
            prev_facts_count = len(self.knowledge_state.key_facts)

            self._update_knowledge(search_results)

            # Step 4: Check if we have a confident answer
            self._assess_answer()

            # Check if we made progress
            confidence_change = (
                self.knowledge_state.confidence - prev_confidence
            )
            new_facts = len(self.knowledge_state.key_facts) - prev_facts_count

            if confidence_change < 0.01 and new_facts == 0:
                # No significant progress - add a finding about this
                self.findings.append(
                    {
                        "phase": f"Low Progress Alert (Step {self.knowledge_state.iteration})",
                        "content": f"Search '{next_search[:50]}...' yielded limited new information. May need to adjust search approach.",
                        "timestamp": self._get_timestamp(),
                    }
                )

                # Add a flag to indicate we need a different approach
                if not hasattr(self.knowledge_state, "low_progress_count"):
                    self.knowledge_state.low_progress_count = 0
                self.knowledge_state.low_progress_count += 1

            if self.progress_callback:
                candidates_count = len(self.knowledge_state.candidate_answers)
                self.progress_callback(
                    f"Step {self.knowledge_state.iteration} complete: {len(self.knowledge_state.key_facts)} facts, {candidates_count} candidates, {self.knowledge_state.confidence:.0%} confident",
                    progress + 10,
                    {
                        "phase": "step_complete",
                        "iteration": self.knowledge_state.iteration,
                        "facts": len(self.knowledge_state.key_facts),
                        "candidates": candidates_count,
                        "confidence": self.knowledge_state.confidence,
                    },
                )

            # Create detailed iteration findings
            iteration_summary = f"""
**Search Query**: {next_search}

**New Facts Discovered**:
{chr(10).join(f"- {fact}" for fact in self.knowledge_state.key_facts[-3:]) if self.knowledge_state.key_facts else "- No new facts in this iteration"}

**Current Candidates** ({len(self.knowledge_state.candidate_answers)} total):
{chr(10).join(f"- {c['answer']} (confidence: {c['confidence']:.0%})" for c in sorted(self.knowledge_state.candidate_answers, key=lambda x: x["confidence"], reverse=True)[:5]) if self.knowledge_state.candidate_answers else "- No candidates yet"}

**Remaining Questions**:
{chr(10).join(f"- {u}" for u in self.knowledge_state.uncertainties[:3]) if self.knowledge_state.uncertainties else "- No specific uncertainties"}

**Progress**: {self.knowledge_state.confidence:.0%} confident after {self.knowledge_state.iteration} steps
"""

            # Add iteration summary to findings
            self.findings.append(
                {
                    "phase": f"Iteration {self.knowledge_state.iteration}",
                    "content": iteration_summary.strip(),
                    "search_query": next_search,
                    "key_findings": self.knowledge_state.key_facts[
                        -3:
                    ],  # Last 3 facts
                    "confidence": self.knowledge_state.confidence,
                    "timestamp": self._get_timestamp(),
                }
            )

        # Final synthesis
        if self.progress_callback:
            self.progress_callback(
                f"Creating final answer based on {len(self.knowledge_state.key_facts)} facts discovered",
                90,
                {
                    "phase": "final_synthesis",
                    "facts_count": len(self.knowledge_state.key_facts),
                    "total_searches": len(self.knowledge_state.search_history),
                    "final_confidence": self.knowledge_state.confidence,
                },
            )

        final_result = self._synthesize_final_answer()

        if self.progress_callback:
            self.progress_callback(
                f"Analysis complete - {self.knowledge_state.iteration} reasoning steps, {len(self.knowledge_state.key_facts)} facts found",
                100,
                {
                    "phase": "complete",
                    "strategy": "iterative_reasoning",
                    "total_iterations": self.knowledge_state.iteration,
                    "facts_discovered": len(self.knowledge_state.key_facts),
                    "final_confidence": self.knowledge_state.confidence,
                },
            )

        return final_result

    def _decide_next_search(self) -> Optional[str]:
        """Decide what to search next based on current knowledge.

        Returns:
            Next search query, or None if done
        """
        # Check for low progress
        low_progress_warning = ""
        if (
            hasattr(self.knowledge_state, "low_progress_count")
            and self.knowledge_state.low_progress_count > 1
        ):
            low_progress_warning = f"\nNOTE: {self.knowledge_state.low_progress_count} recent searches yielded limited new information. Try a significantly different search approach."

        prompt = f"""Based on our current knowledge, decide what to search next.

{self.knowledge_state.to_string()}{low_progress_warning}

Consider:
1. What specific information would help answer the original query?
2. What uncertainties should we resolve?
3. Should we verify any candidate answers?
4. Do we need more specific or broader information?
5. Can we combine multiple constraints into a more targeted search?
6. Are there multiple candidates with similar confidence? If so, what searches would help distinguish between them?
7. If recent searches haven't been productive, what completely different approach could we try?

For puzzle-like queries with specific clues, try to:
- Search for locations that match multiple criteria at once
- Use specific place names when possible
- Include relevant statistics or dates mentioned
- If multiple candidates exist, search for distinguishing features
- If searches are repetitive, try broader regional searches or different constraint combinations

If we have one clear candidate with high confidence and low uncertainty, respond with "DONE".
If multiple candidates have similar confidence, continue searching to distinguish between them.

Otherwise, provide:
NEXT_SEARCH: [specific search query that targets the constraints]
REASONING: [why this search will help]
EXPECTED_OUTCOME: [what we hope to learn]
"""

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

        logger.debug(f"LLM response for next search: {content[:200]}...")

        if "DONE" in content:
            logger.info("LLM decided no more searches needed")
            return None

        # Parse response
        lines = content.strip().split("\n")
        next_search = None
        reasoning = ""

        for line in lines:
            if line.startswith("NEXT_SEARCH:"):
                next_search = line.split(":", 1)[1].strip()
            elif line.startswith("REASONING:"):
                reasoning = line.split(":", 1)[1].strip()

        if next_search:
            logger.info(f"Next search: {next_search}")
            logger.info(f"Reasoning: {reasoning}")
        else:
            logger.warning("Could not extract next search from LLM response")
            # Try to extract any search-like query from the response
            for line in lines:
                if (
                    len(line) > 10 and "?" not in line
                ):  # Heuristic for search queries
                    next_search = line.strip()
                    logger.info(f"Extracted fallback search: {next_search}")
                    break

        return next_search

    def _execute_search(self, search_query: str) -> Dict:
        """Execute a search using the source-based strategy.

        Args:
            search_query: The query to search

        Returns:
            Search results
        """
        logger.info(f"Executing search: {search_query}")

        # Use source-based strategy for actual search
        source_strategy = SourceBasedSearchStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            include_text_content=True,
            use_cross_engine_filter=True,
            use_atomic_facts=True,
        )

        source_strategy.max_iterations = self.search_iterations_per_round
        source_strategy.questions_per_iteration = self.questions_per_search

        if self.progress_callback:
            # Create a wrapped callback that includes our iteration info
            def wrapped_callback(message, progress, data):
                data["parent_iteration"] = self.knowledge_state.iteration
                data["parent_strategy"] = "iterative_reasoning"
                # Make the message clearer for the GUI
                if "Searching" in message:
                    display_message = (
                        f"Step {self.knowledge_state.iteration}: {message}"
                    )
                elif "questions" in message.lower():
                    display_message = f"Step {self.knowledge_state.iteration}: Generating follow-up questions"
                elif "findings" in message.lower():
                    display_message = f"Step {self.knowledge_state.iteration}: Processing findings"
                else:
                    display_message = (
                        f"Step {self.knowledge_state.iteration}: {message}"
                    )

                self.progress_callback(display_message, progress, data)

            source_strategy.set_progress_callback(wrapped_callback)

        results = source_strategy.analyze_topic(search_query)

        # Track search history
        self.knowledge_state.search_history.append(
            {
                "query": search_query,
                "timestamp": self._get_timestamp(),
                "findings_count": len(results.get("findings", [])),
                "sources_count": len(results.get("sources", [])),
            }
        )

        # Update our tracking
        if "questions_by_iteration" in results:
            self.questions_by_iteration.extend(
                results["questions_by_iteration"]
            )

        return results

    def _update_knowledge(self, search_results: Dict):
        """Update knowledge state with new findings.

        Args:
            search_results: Results from the search
        """
        current_knowledge = search_results.get("current_knowledge", "")

        if not current_knowledge:
            return

        prompt = f"""Update our knowledge based on new search results.

Current Knowledge State:
{self.knowledge_state.to_string()}

New Information:
{current_knowledge[:2000]}  # Truncate if too long

Extract:
1. KEY_FACTS: Specific facts that help answer the original query
2. NEW_UNCERTAINTIES: Questions that arose from this search
3. CANDIDATE_ANSWERS: Possible answers to the original query with confidence (0-1)
4. RESOLVED_UNCERTAINTIES: Which previous uncertainties were resolved

Format:
KEY_FACTS:
- [fact 1]
- [fact 2]

NEW_UNCERTAINTIES:
- [uncertainty 1]

CANDIDATE_ANSWERS:
- [answer]: [confidence]

RESOLVED_UNCERTAINTIES:
- [resolved uncertainty]
"""

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

        # Parse response and update knowledge state
        lines = content.strip().split("\n")
        current_section = None

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if line.startswith("KEY_FACTS:"):
                current_section = "facts"
            elif line.startswith("NEW_UNCERTAINTIES:"):
                current_section = "uncertainties"
            elif line.startswith("CANDIDATE_ANSWERS:"):
                current_section = "candidates"
            elif line.startswith("RESOLVED_UNCERTAINTIES:"):
                current_section = "resolved"
            elif line.startswith("-") and current_section:
                item = line[1:].strip()

                if current_section == "facts":
                    if item not in self.knowledge_state.key_facts:
                        self.knowledge_state.key_facts.append(item)
                elif current_section == "uncertainties":
                    if item not in self.knowledge_state.uncertainties:
                        self.knowledge_state.uncertainties.append(item)
                elif current_section == "candidates" and ":" in item:
                    answer, confidence_str = item.split(":", 1)
                    answer = answer.strip()

                    # Skip invalid answers
                    if (
                        answer.lower() in ["none", "unknown", "n/a", ""]
                        or len(answer) < 3
                    ):
                        continue

                    try:
                        confidence = float(confidence_str.strip())
                    except ValueError:
                        confidence = 0.5

                    # Clean up answer text (remove duplicated confidence info)
                    if "(confidence" in answer:
                        answer = answer.split("(confidence")[0].strip()

                    # Update or add candidate
                    existing = False
                    for candidate in self.knowledge_state.candidate_answers:
                        if candidate["answer"].lower() == answer.lower():
                            candidate["confidence"] = max(
                                candidate["confidence"], confidence
                            )
                            existing = True
                            break

                    if not existing:
                        self.knowledge_state.candidate_answers.append(
                            {"answer": answer, "confidence": confidence}
                        )
                elif current_section == "resolved":
                    # Remove from uncertainties if present
                    self.knowledge_state.uncertainties = [
                        u
                        for u in self.knowledge_state.uncertainties
                        if item.lower() not in u.lower()
                    ]

        logger.info(
            f"Knowledge updated: {len(self.knowledge_state.key_facts)} facts, "
            f"{len(self.knowledge_state.uncertainties)} uncertainties"
        )

    def _assess_answer(self):
        """Assess our confidence in the current answer."""
        if not self.knowledge_state.candidate_answers:
            self.knowledge_state.confidence = 0.0
            return

        # Sort candidates by confidence
        sorted_candidates = sorted(
            self.knowledge_state.candidate_answers,
            key=lambda x: x["confidence"],
            reverse=True,
        )

        best_candidate = sorted_candidates[0]
        base_confidence = best_candidate["confidence"]

        # Check if multiple candidates have similar confidence
        confidence_spread = 0.0
        if len(sorted_candidates) > 1:
            second_best = sorted_candidates[1]
            confidence_spread = (
                best_candidate["confidence"] - second_best["confidence"]
            )

            # If top two candidates are very close, reduce overall confidence
            if confidence_spread < 0.1:  # Less than 10% difference
                base_confidence *= 0.8  # Reduce confidence by 20%
                logger.info(
                    f"Multiple candidates with similar confidence: {best_candidate['answer']} ({best_candidate['confidence']:.0%}) vs {second_best['answer']} ({second_best['confidence']:.0%})"
                )

        # Consider multiple factors for overall confidence
        # Boost confidence if we have supporting facts
        fact_boost = min(len(self.knowledge_state.key_facts) * 0.05, 0.2)

        # Reduce confidence if we have many uncertainties
        uncertainty_penalty = min(
            len(self.knowledge_state.uncertainties) * 0.05, 0.2
        )

        # Boost if we've done multiple searches that confirm
        search_boost = min(len(self.knowledge_state.search_history) * 0.02, 0.1)

        # Penalty for too many candidates
        candidate_penalty = 0.0
        if len(self.knowledge_state.candidate_answers) > 3:
            candidate_penalty = min(
                (len(self.knowledge_state.candidate_answers) - 3) * 0.05, 0.15
            )

        self.knowledge_state.confidence = min(
            base_confidence
            + fact_boost
            + search_boost
            - uncertainty_penalty
            - candidate_penalty,
            0.95,  # Cap at 95% to leave room for synthesis
        )

        logger.info(
            f"Current confidence: {self.knowledge_state.confidence:.1%} (best: {best_candidate['answer']} at {best_candidate['confidence']:.0%}, spread: {confidence_spread:.1%})"
        )

    def _synthesize_final_answer(self) -> Dict:
        """Synthesize final answer from accumulated knowledge.

        Returns:
            Final results dictionary
        """
        # Get best answer (removed - not used anymore)

        # Create comprehensive answer
        prompt = f"""Provide a final answer to the original query based on our research.

{self.knowledge_state.to_string()}

Requirements:
1. Directly answer the original query
2. Reference the key facts that support the answer
3. Acknowledge any remaining uncertainties
4. Be concise but complete

Final Answer:"""

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

        # Add comprehensive final synthesis to findings
        synthesis_summary = f"""
**Final Answer**:
{final_answer}

**Research Summary**:
- Completed {self.knowledge_state.iteration} reasoning steps
- Discovered {len(self.knowledge_state.key_facts)} key facts
- Evaluated {len(self.knowledge_state.candidate_answers)} candidate answers
- Final confidence: {self.knowledge_state.confidence:.0%}

**Key Facts That Led to This Answer**:
{chr(10).join(f"{i + 1}. {fact}" for i, fact in enumerate(self.knowledge_state.key_facts[:10]))}

**Search History**:
{chr(10).join(f"{i + 1}. {search['query']}" for i, search in enumerate(self.knowledge_state.search_history))}
""".strip()

        self.findings.append(
            {
                "phase": "Final Synthesis",
                "content": synthesis_summary,
                "confidence": self.knowledge_state.confidence,
                "iterations": self.knowledge_state.iteration,
                "timestamp": self._get_timestamp(),
            }
        )

        # Compile all questions
        questions_dict = {}
        for i, questions in enumerate(self.questions_by_iteration):
            if isinstance(questions, list):
                questions_dict[i + 1] = questions
            elif isinstance(questions, dict):
                questions_dict.update(questions)

        # Format findings
        formatted_findings = format_findings(
            self.findings, final_answer, questions_dict
        )

        return {
            "current_knowledge": final_answer,
            "formatted_findings": formatted_findings,
            "findings": self.findings,
            "iterations": self.knowledge_state.iteration,
            "questions_by_iteration": questions_dict,
            "all_links_of_system": self.all_links_of_system,
            "sources": self.all_links_of_system,
            "knowledge_state": {
                "key_facts": self.knowledge_state.key_facts,
                "uncertainties": self.knowledge_state.uncertainties,
                "candidate_answers": self.knowledge_state.candidate_answers,
                "confidence": self.knowledge_state.confidence,
                "search_history": self.knowledge_state.search_history,
            },
            "strategy": "iterative_reasoning",
            "questions": {
                "total": sum(
                    len(q) if isinstance(q, list) else 1
                    for q in questions_dict.values()
                ),
                "by_iteration": questions_dict,
            },
        }

    def _get_timestamp(self) -> str:
        """Get current timestamp for findings."""
        return datetime.now(UTC).isoformat()
