"""
RNSR Self-Reflection Loop

Implements iterative self-correction where the system:
1. Generates an initial answer
2. Critiques its own answer
3. If issues found, re-navigates with critique as context
4. Repeats until confident or max iterations

Based on self-reflection patterns from:
- Reflexion (Shinn et al.)
- Self-Refine (Madaan et al.)
- Constitutional AI principles

Key insight: LLMs can often identify problems in their own outputs
that they couldn't avoid in initial generation.
"""

from __future__ import annotations

import json
import re
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable

import structlog

logger = structlog.get_logger(__name__)


# =============================================================================
# Self-Reflection Prompts
# =============================================================================

CRITIQUE_PROMPT = """You are a critical reviewer. Analyze this answer for potential issues.

QUESTION: {question}

ANSWER: {answer}

EVIDENCE USED:
{evidence}

Critically evaluate:
1. ACCURACY: Does the evidence actually support this answer?
2. COMPLETENESS: Is anything important missing?
3. CONTRADICTIONS: Does any evidence contradict the answer?
4. SPECIFICITY: Is the answer too vague or too specific?
5. ASSUMPTIONS: Are there unstated assumptions?

If you find issues, explain them clearly.
If the answer is good, say "NO ISSUES FOUND".

Respond in JSON:
{{
    "has_issues": true/false,
    "issues": [
        {{"type": "accuracy|completeness|contradiction|specificity|assumption", "description": "...", "severity": "high|medium|low"}}
    ],
    "confidence_in_critique": 0.0-1.0,
    "suggested_improvements": ["..."],
    "should_retry": true/false
}}"""


REFINEMENT_PROMPT = """You are improving an answer based on feedback.

ORIGINAL QUESTION: {question}

PREVIOUS ANSWER: {previous_answer}

CRITIQUE/ISSUES FOUND:
{critique}

EVIDENCE AVAILABLE:
{evidence}

Generate an IMPROVED answer that addresses the identified issues.
Be specific and directly address each criticism.

Respond with ONLY the improved answer, no meta-commentary."""


VERIFICATION_PROMPT = """Compare these two answers and determine which is better.

QUESTION: {question}

ANSWER A (Original):
{answer_a}

ANSWER B (Refined):
{answer_b}

Which answer is:
1. More accurate?
2. More complete?
3. Better supported by evidence?

Respond in JSON:
{{
    "better_answer": "A" or "B",
    "confidence": 0.0-1.0,
    "reasoning": "..."
}}"""


# =============================================================================
# Data Models
# =============================================================================

class IssueType(str, Enum):
    """Types of issues that can be identified."""
    
    ACCURACY = "accuracy"
    COMPLETENESS = "completeness"
    CONTRADICTION = "contradiction"
    SPECIFICITY = "specificity"
    ASSUMPTION = "assumption"
    HALLUCINATION = "hallucination"


class IssueSeverity(str, Enum):
    """Severity of identified issues."""
    
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


@dataclass
class Issue:
    """An issue identified during self-critique."""
    
    type: IssueType
    description: str
    severity: IssueSeverity
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "type": self.type.value,
            "description": self.description,
            "severity": self.severity.value,
        }


@dataclass
class CritiqueResult:
    """Result of self-critique."""
    
    has_issues: bool = False
    issues: list[Issue] = field(default_factory=list)
    confidence: float = 0.5
    suggested_improvements: list[str] = field(default_factory=list)
    should_retry: bool = False
    raw_response: str = ""
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "has_issues": self.has_issues,
            "issues": [i.to_dict() for i in self.issues],
            "confidence": self.confidence,
            "suggested_improvements": self.suggested_improvements,
            "should_retry": self.should_retry,
        }


@dataclass
class ReflectionIteration:
    """One iteration of the reflection loop."""
    
    iteration: int
    answer: str
    critique: CritiqueResult | None = None
    improved_answer: str | None = None
    improvement_accepted: bool = False
    duration_ms: float = 0.0
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "iteration": self.iteration,
            "answer": self.answer,
            "critique": self.critique.to_dict() if self.critique else None,
            "improved_answer": self.improved_answer,
            "improvement_accepted": self.improvement_accepted,
            "duration_ms": self.duration_ms,
        }


@dataclass
class ReflectionResult:
    """Complete result of self-reflection process."""
    
    original_answer: str = ""
    final_answer: str = ""
    question: str = ""
    
    # Iterations
    iterations: list[ReflectionIteration] = field(default_factory=list)
    total_iterations: int = 0
    
    # Outcome
    improved: bool = False
    final_confidence: float = 0.0
    all_issues: list[Issue] = field(default_factory=list)
    
    # Timing
    total_duration_ms: float = 0.0
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "original_answer": self.original_answer,
            "final_answer": self.final_answer,
            "question": self.question,
            "iterations": [i.to_dict() for i in self.iterations],
            "total_iterations": self.total_iterations,
            "improved": self.improved,
            "final_confidence": self.final_confidence,
            "all_issues": [i.to_dict() for i in self.all_issues],
            "total_duration_ms": self.total_duration_ms,
        }


# =============================================================================
# Self-Reflection Engine
# =============================================================================


class SelfReflectionEngine:
    """
    Implements iterative self-correction for answers.
    
    Flow:
    1. Take initial answer
    2. Generate critique (what could be wrong?)
    3. If issues found, generate improved answer
    4. Verify improvement is actually better
    5. Repeat until confident or max iterations
    """
    
    def __init__(
        self,
        llm_fn: Callable[[str], str] | None = None,
        max_iterations: int = 3,
        min_confidence_threshold: float = 0.8,
        accept_improvement_threshold: float = 0.6,
        enable_verification: bool = True,
    ):
        """
        Initialize the self-reflection engine.
        
        Args:
            llm_fn: LLM function for critique and refinement.
            max_iterations: Maximum reflection iterations.
            min_confidence_threshold: Stop if critique confidence exceeds this.
            accept_improvement_threshold: Accept improvement if confidence exceeds this.
            enable_verification: Verify improvements are actually better.
        """
        self.llm_fn = llm_fn
        self.max_iterations = max_iterations
        self.min_confidence_threshold = min_confidence_threshold
        self.accept_improvement_threshold = accept_improvement_threshold
        self.enable_verification = enable_verification
        
        # Learning: track which issues are commonly found
        self._issue_stats: dict[str, int] = {}
    
    def set_llm_function(self, llm_fn: Callable[[str], str]) -> None:
        """Set the LLM function."""
        self.llm_fn = llm_fn
    
    def reflect(
        self,
        answer: str,
        question: str,
        evidence: str = "",
        navigate_fn: Callable[[str], str] | None = None,
    ) -> ReflectionResult:
        """
        Perform self-reflection on an answer.
        
        Args:
            answer: The initial answer to reflect on.
            question: The original question.
            evidence: Evidence that was used to generate the answer.
            navigate_fn: Optional function to re-navigate with new context.
            
        Returns:
            ReflectionResult with final answer and iteration history.
        """
        if self.llm_fn is None:
            logger.warning("no_llm_configured_for_reflection")
            return ReflectionResult(
                original_answer=answer,
                final_answer=answer,
                question=question,
            )
        
        start_time = time.time()
        
        result = ReflectionResult(
            original_answer=answer,
            final_answer=answer,
            question=question,
        )
        
        current_answer = answer
        
        for iteration in range(self.max_iterations):
            iter_start = time.time()
            
            logger.info(
                "reflection_iteration_started",
                iteration=iteration + 1,
                max=self.max_iterations,
            )
            
            # Step 1: Critique the current answer
            critique = self._critique(current_answer, question, evidence)
            
            iter_result = ReflectionIteration(
                iteration=iteration + 1,
                answer=current_answer,
                critique=critique,
            )
            
            # Collect issues for statistics
            for issue in critique.issues:
                self._issue_stats[issue.type.value] = \
                    self._issue_stats.get(issue.type.value, 0) + 1
                result.all_issues.append(issue)
            
            # Check if we should stop
            if not critique.has_issues or not critique.should_retry:
                logger.info(
                    "reflection_no_issues",
                    iteration=iteration + 1,
                    confidence=critique.confidence,
                )
                iter_result.duration_ms = (time.time() - iter_start) * 1000
                result.iterations.append(iter_result)
                break
            
            # Step 2: Generate improved answer
            improved_answer = self._refine(
                current_answer, question, critique, evidence
            )
            
            iter_result.improved_answer = improved_answer
            
            # Step 3: Verify improvement (optional)
            if self.enable_verification and improved_answer:
                is_better = self._verify_improvement(
                    question, current_answer, improved_answer
                )
                iter_result.improvement_accepted = is_better
                
                if is_better:
                    current_answer = improved_answer
                    result.improved = True
                    logger.info(
                        "improvement_accepted",
                        iteration=iteration + 1,
                    )
                else:
                    logger.info(
                        "improvement_rejected",
                        iteration=iteration + 1,
                    )
            elif improved_answer:
                # Accept without verification
                current_answer = improved_answer
                iter_result.improvement_accepted = True
                result.improved = True
            
            iter_result.duration_ms = (time.time() - iter_start) * 1000
            result.iterations.append(iter_result)
            
            # Check confidence threshold
            if critique.confidence >= self.min_confidence_threshold:
                logger.info(
                    "confidence_threshold_reached",
                    confidence=critique.confidence,
                    threshold=self.min_confidence_threshold,
                )
                break
        
        result.final_answer = current_answer
        result.total_iterations = len(result.iterations)
        result.total_duration_ms = (time.time() - start_time) * 1000
        
        # Calculate final confidence
        if result.iterations and result.iterations[-1].critique:
            last_critique = result.iterations[-1].critique
            result.final_confidence = last_critique.confidence if not last_critique.has_issues else 1.0 - (len(last_critique.issues) * 0.1)
        else:
            result.final_confidence = 0.7  # Default
        
        logger.info(
            "reflection_complete",
            iterations=result.total_iterations,
            improved=result.improved,
            final_confidence=result.final_confidence,
            duration_ms=result.total_duration_ms,
        )
        
        return result
    
    def _critique(
        self,
        answer: str,
        question: str,
        evidence: str,
    ) -> CritiqueResult:
        """Generate a critique of the answer."""
        prompt = CRITIQUE_PROMPT.format(
            question=question,
            answer=answer,
            evidence=evidence[:2000] if evidence else "No specific evidence provided.",
        )
        
        try:
            response = self.llm_fn(prompt)
            return self._parse_critique(response)
            
        except Exception as e:
            logger.warning("critique_failed", error=str(e))
            return CritiqueResult(
                has_issues=False,
                confidence=0.5,
                raw_response=str(e),
            )
    
    def _parse_critique(self, response: str) -> CritiqueResult:
        """Parse critique response into structured format."""
        result = CritiqueResult(raw_response=response)
        
        # Check for "NO ISSUES FOUND"
        if "NO ISSUES FOUND" in response.upper():
            result.has_issues = False
            result.confidence = 0.9
            return result
        
        # Parse JSON
        try:
            json_match = re.search(r'\{[\s\S]*\}', response)
            if not json_match:
                return result
            
            data = json.loads(json_match.group())
            
            result.has_issues = data.get("has_issues", False)
            result.confidence = data.get("confidence_in_critique", 0.5)
            result.should_retry = data.get("should_retry", False)
            result.suggested_improvements = data.get("suggested_improvements", [])
            
            for issue_data in data.get("issues", []):
                try:
                    issue = Issue(
                        type=IssueType(issue_data.get("type", "accuracy")),
                        description=issue_data.get("description", ""),
                        severity=IssueSeverity(issue_data.get("severity", "medium")),
                    )
                    result.issues.append(issue)
                except ValueError:
                    pass
            
        except json.JSONDecodeError:
            # If JSON parsing fails, look for issue indicators
            if any(word in response.lower() for word in ["issue", "problem", "incorrect", "missing"]):
                result.has_issues = True
                result.should_retry = True
        
        return result
    
    def _refine(
        self,
        answer: str,
        question: str,
        critique: CritiqueResult,
        evidence: str,
    ) -> str:
        """Generate an improved answer based on critique."""
        # Format critique for prompt
        critique_text = []
        for issue in critique.issues:
            critique_text.append(f"- [{issue.severity.value.upper()}] {issue.type.value}: {issue.description}")
        
        if critique.suggested_improvements:
            critique_text.append("\nSuggested improvements:")
            for suggestion in critique.suggested_improvements:
                critique_text.append(f"- {suggestion}")
        
        prompt = REFINEMENT_PROMPT.format(
            question=question,
            previous_answer=answer,
            critique="\n".join(critique_text) if critique_text else "No specific issues identified.",
            evidence=evidence[:2000] if evidence else "Use your knowledge to improve the answer.",
        )
        
        try:
            response = self.llm_fn(prompt)
            return response.strip()
            
        except Exception as e:
            logger.warning("refinement_failed", error=str(e))
            return ""
    
    def _verify_improvement(
        self,
        question: str,
        original: str,
        improved: str,
    ) -> bool:
        """Verify that the improved answer is actually better."""
        if not improved:
            return False
        
        prompt = VERIFICATION_PROMPT.format(
            question=question,
            answer_a=original,
            answer_b=improved,
        )
        
        try:
            response = self.llm_fn(prompt)
            
            # Parse response
            json_match = re.search(r'\{[\s\S]*\}', response)
            if not json_match:
                # Default to accepting improvement
                return True
            
            data = json.loads(json_match.group())
            better = data.get("better_answer", "B")
            confidence = data.get("confidence", 0.5)
            
            # Accept B (improved) if confident enough
            return better == "B" and confidence >= self.accept_improvement_threshold
            
        except Exception as e:
            logger.warning("verification_failed", error=str(e))
            # Default to accepting improvement
            return True
    
    def get_issue_stats(self) -> dict[str, int]:
        """Get statistics on issues found across all reflections."""
        return dict(self._issue_stats)


# =============================================================================
# Convenience Functions
# =============================================================================


def reflect_on_answer(
    answer: str,
    question: str,
    evidence: str = "",
    llm_fn: Callable[[str], str] | None = None,
    max_iterations: int = 2,
) -> ReflectionResult:
    """
    Perform self-reflection on an answer.
    
    Simple interface for one-off reflection.
    
    Args:
        answer: The answer to reflect on.
        question: The original question.
        evidence: Evidence used.
        llm_fn: LLM function (uses default if not provided).
        max_iterations: Maximum iterations.
        
    Returns:
        ReflectionResult with final answer.
    """
    if llm_fn is None:
        try:
            from rnsr.llm import get_llm
            llm = get_llm()
            llm_fn = lambda p: str(llm.complete(p))
        except Exception as e:
            logger.warning("no_llm_available", error=str(e))
            return ReflectionResult(
                original_answer=answer,
                final_answer=answer,
                question=question,
            )
    
    engine = SelfReflectionEngine(
        llm_fn=llm_fn,
        max_iterations=max_iterations,
    )
    
    return engine.reflect(answer, question, evidence)
