"""
R-Zero challenger-solver self-refinement loops.

Implements an AlphaZero-inspired pattern where a solver produces solutions
and a challenger critiques them with improvement suggestions.

Requirements: SEC-07
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import List, Optional, Callable, Any, Protocol
from enum import Enum
import asyncio
import inspect

# Optional DSPy import with graceful fallback
try:
    import dspy

    HAS_DSPY = True
except ImportError:
    HAS_DSPY = False
    dspy = None


def _utc_now_iso() -> str:
    """Get current UTC time as ISO format string."""
    return datetime.now(timezone.utc).isoformat()


class RefinementStatus(Enum):
    """Status of R-Zero refinement loop."""

    IN_PROGRESS = "in_progress"
    THRESHOLD_MET = "threshold_met"
    MAX_ROUNDS_REACHED = "max_rounds_reached"
    ERROR = "error"


@dataclass
class ChallengerFeedback:
    """
    Feedback from challenger to solver.

    Captures critique of a solution and suggestions for improvement,
    enabling the self-refinement loop.
    """

    critique: str
    alternative_approach: str
    suggested_improvements: List[str] = field(default_factory=list)
    confidence: float = 0.5
    timestamp: str = field(default_factory=_utc_now_iso)

    def to_dict(self) -> dict:
        """
        Convert feedback to dictionary for serialization.

        Returns:
            Dictionary representation suitable for JSON storage.
        """
        return {
            "critique": self.critique,
            "alternative_approach": self.alternative_approach,
            "suggested_improvements": self.suggested_improvements,
            "confidence": self.confidence,
            "timestamp": self.timestamp,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "ChallengerFeedback":
        """
        Create ChallengerFeedback from dictionary.

        Args:
            data: Dictionary with feedback fields.

        Returns:
            ChallengerFeedback instance.
        """
        return cls(
            critique=data.get("critique", ""),
            alternative_approach=data.get("alternative_approach", ""),
            suggested_improvements=data.get("suggested_improvements", []),
            confidence=data.get("confidence", 0.5),
            timestamp=data.get("timestamp", _utc_now_iso()),
        )

    def __post_init__(self):
        """Validate and normalize fields after initialization."""
        # Ensure suggested_improvements is always a list
        if self.suggested_improvements is None:
            self.suggested_improvements = []

        # Clamp confidence to valid range
        self.confidence = max(0.0, min(1.0, self.confidence))

        # Ensure timestamp has a value
        if not self.timestamp:
            self.timestamp = _utc_now_iso()


# DSPy Signature for challenger (defined only if DSPy available)
if HAS_DSPY:

    class ChallengerSignature(dspy.Signature):
        """Critique a solution and propose improvements."""

        solution: str = dspy.InputField(desc="Current solution to critique")
        context: str = dspy.InputField(desc="Task context and requirements")
        reward: float = dspy.InputField(desc="Current reward score (0-1)")
        round_num: int = dspy.InputField(desc="Current refinement round")

        critique: str = dspy.OutputField(desc="Analysis of solution weaknesses")
        alternative_approach: str = dspy.OutputField(
            desc="Alternative solution approach"
        )
        improvements: List[str] = dspy.OutputField(
            desc="Specific improvements to apply"
        )
else:
    ChallengerSignature = None


class RewardFunction(Protocol):
    """Protocol for reward functions used in R-Zero loops."""

    def __call__(self, input_data: Any, prediction: Any, context: str = "") -> float:
        """
        Evaluate a prediction and return a reward score.

        Args:
            input_data: The original task input.
            prediction: The solver's output/prediction.
            context: Additional context for evaluation.

        Returns:
            Reward score between 0.0 and 1.0.
        """
        ...


@dataclass
class RefinementResult:
    """
    Result of R-Zero refinement loop.

    Captures the final solution, improvement trajectory, and all
    feedback received during the refinement process.
    """

    final_solution: Any
    final_reward: float
    status: RefinementStatus
    rounds_completed: int
    improvement_history: List[float] = field(default_factory=list)
    feedback_history: List[ChallengerFeedback] = field(default_factory=list)
    timestamp: str = field(default_factory=_utc_now_iso)
    error_message: Optional[str] = None

    @property
    def improved(self) -> bool:
        """Check if refinement improved the solution."""
        if len(self.improvement_history) < 2:
            return False
        return self.improvement_history[-1] > self.improvement_history[0]

    def to_dict(self) -> dict:
        """
        Convert result to dictionary for serialization.

        Returns:
            Dictionary representation suitable for JSON storage.
        """
        return {
            "final_solution": self.final_solution,
            "final_reward": self.final_reward,
            "status": self.status.value,
            "rounds_completed": self.rounds_completed,
            "improvement_history": self.improvement_history,
            "feedback_history": [fb.to_dict() for fb in self.feedback_history],
            "timestamp": self.timestamp,
            "error_message": self.error_message,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "RefinementResult":
        """
        Create RefinementResult from dictionary.

        Args:
            data: Dictionary with result fields.

        Returns:
            RefinementResult instance.
        """
        feedback_history = [
            ChallengerFeedback.from_dict(fb) for fb in data.get("feedback_history", [])
        ]

        return cls(
            final_solution=data.get("final_solution"),
            final_reward=data.get("final_reward", 0.0),
            status=RefinementStatus(data.get("status", "in_progress")),
            rounds_completed=data.get("rounds_completed", 0),
            improvement_history=data.get("improvement_history", []),
            feedback_history=feedback_history,
            timestamp=data.get("timestamp", _utc_now_iso()),
            error_message=data.get("error_message"),
        )


class RZeroLoop:
    """
    Challenger-solver loop for self-refinement.

    Implements an AlphaZero-inspired pattern where:
    1. A solver produces candidate solutions
    2. A challenger critiques them with improvement suggestions
    3. Feedback is incorporated in subsequent solver calls
    4. The loop continues until quality threshold is met or max rounds reached
    """

    DEFAULT_MAX_ROUNDS = 3
    DEFAULT_THRESHOLD = 0.9

    def __init__(
        self,
        solver: Any,  # Callable or dspy.Module
        challenger: Optional[Any] = None,
        reward_fn: Optional[RewardFunction] = None,
        max_rounds: int = DEFAULT_MAX_ROUNDS,
        threshold: float = DEFAULT_THRESHOLD,
    ):
        """
        Initialize the R-Zero refinement loop.

        Args:
            solver: A callable or dspy.Module that produces solutions.
            challenger: A callable or dspy.Module that critiques solutions.
                       If None, refinement happens without challenger feedback.
            reward_fn: Function to evaluate solution quality (0.0-1.0).
                      If None, uses default reward function.
            max_rounds: Maximum number of refinement rounds.
            threshold: Reward threshold to stop refinement early.
        """
        self.solver = solver
        self.challenger = challenger
        self.reward_fn = reward_fn or self._default_reward
        self.max_rounds = max_rounds
        self.threshold = threshold

    def _default_reward(
        self, input_data: Any, prediction: Any, context: str = ""
    ) -> float:
        """
        Default reward function if none provided.

        Provides basic reward based on prediction presence and structure.
        """
        if prediction is None:
            return 0.0
        if isinstance(prediction, dict):
            return prediction.get("reward", 0.5)
        if isinstance(prediction, str) and len(prediction) > 0:
            return 0.5
        return 0.5

    async def execute(
        self,
        task: str,
        context: str = "",
        initial_hint: Optional[str] = None,
    ) -> RefinementResult:
        """
        Execute the challenger-solver loop.

        1. Call solver with task and context
        2. Evaluate with reward_fn
        3. If reward >= threshold, return immediately
        4. Otherwise, call challenger for feedback
        5. Incorporate feedback (as hint) in next solver call
        6. Track improvement history
        7. Return best result found

        Args:
            task: The task description to solve.
            context: Additional context for the task.
            initial_hint: Optional hint for first solver call.

        Returns:
            RefinementResult with final solution and improvement history.
        """
        improvement_history: List[float] = []
        feedback_history: List[ChallengerFeedback] = []
        best_solution = None
        best_reward = 0.0
        current_feedback: Optional[ChallengerFeedback] = None
        current_hint: Optional[str] = initial_hint
        error_message: Optional[str] = None

        for round_num in range(self.max_rounds):
            try:
                # Call solver with current hint string
                solution = await self._call_solver(task, context, current_hint)

                # Evaluate solution
                reward = self.reward_fn(task, solution, context)
                improvement_history.append(reward)

                # Track best solution
                if reward > best_reward:
                    best_reward = reward
                    best_solution = solution

                # Check threshold
                if reward >= self.threshold:
                    return RefinementResult(
                        final_solution=solution,
                        final_reward=reward,
                        status=RefinementStatus.THRESHOLD_MET,
                        rounds_completed=round_num + 1,
                        improvement_history=improvement_history,
                        feedback_history=feedback_history,
                    )

                # Get challenger feedback if available
                if self.challenger is not None:
                    feedback = await self._call_challenger(
                        solution, context, reward, round_num + 1
                    )
                    if feedback:
                        feedback_history.append(feedback)
                        # Use feedback as hint for next round
                        current_hint = self._format_hint(feedback)

            except Exception as e:
                error_message = str(e)
                # Return best result so far on error
                if best_solution is not None:
                    return RefinementResult(
                        final_solution=best_solution,
                        final_reward=best_reward,
                        status=RefinementStatus.ERROR,
                        rounds_completed=round_num + 1,
                        improvement_history=improvement_history,
                        feedback_history=feedback_history,
                        error_message=error_message,
                    )
                # No solution at all
                return RefinementResult(
                    final_solution=None,
                    final_reward=0.0,
                    status=RefinementStatus.ERROR,
                    rounds_completed=round_num + 1,
                    improvement_history=improvement_history,
                    feedback_history=feedback_history,
                    error_message=error_message,
                )

        # Max rounds reached, return best solution found
        return RefinementResult(
            final_solution=best_solution,
            final_reward=best_reward,
            status=RefinementStatus.MAX_ROUNDS_REACHED,
            rounds_completed=self.max_rounds,
            improvement_history=improvement_history,
            feedback_history=feedback_history,
        )

    async def _call_solver(
        self,
        task: str,
        context: str,
        hint: Optional[str] = None,
    ) -> Any:
        """
        Call the solver to produce a solution.

        Args:
            task: The task to solve.
            context: Additional context.
            hint: Optional hint string from previous round's feedback.

        Returns:
            Solver's prediction/solution.
        """
        # Handle different solver types
        if inspect.iscoroutinefunction(self.solver):
            # Async callable
            if hint:
                return await self.solver(task, context=context, hint=hint)
            return await self.solver(task, context=context)
        elif callable(self.solver):
            # Sync callable
            if hint:
                return self.solver(task, context=context, hint=hint)
            return self.solver(task, context=context)
        elif HAS_DSPY and isinstance(self.solver, dspy.Module):
            # DSPy module - call with context
            try:
                # Run sync DSPy in thread pool
                loop = asyncio.get_event_loop()
                if hint:
                    return await loop.run_in_executor(
                        None, lambda: self.solver(task=task, context=context, hint=hint)
                    )
                return await loop.run_in_executor(
                    None, lambda: self.solver(task=task, context=context)
                )
            except Exception:
                # Fallback if executor fails
                if hint:
                    return self.solver(task=task, context=context, hint=hint)
                return self.solver(task=task, context=context)
        else:
            raise TypeError(f"Unsupported solver type: {type(self.solver)}")

    async def _call_challenger(
        self,
        solution: Any,
        context: str,
        reward: float,
        round_num: int,
    ) -> Optional[ChallengerFeedback]:
        """
        Call the challenger to critique a solution.

        Args:
            solution: The solution to critique.
            context: Task context.
            reward: Current reward score.
            round_num: Current refinement round.

        Returns:
            ChallengerFeedback if available, None otherwise.
        """
        # Convert solution to string for challenger
        solution_str = str(solution) if not isinstance(solution, str) else solution

        # Handle different challenger types
        try:
            if inspect.iscoroutinefunction(self.challenger):
                # Async callable
                result = await self.challenger(
                    solution=solution_str,
                    context=context,
                    reward=reward,
                    round_num=round_num,
                )
            elif callable(self.challenger):
                # Sync callable
                result = self.challenger(
                    solution=solution_str,
                    context=context,
                    reward=reward,
                    round_num=round_num,
                )
            elif HAS_DSPY and isinstance(self.challenger, dspy.Module):
                # DSPy module using ChallengerSignature
                loop = asyncio.get_event_loop()
                result = await loop.run_in_executor(
                    None,
                    lambda: self.challenger(
                        solution=solution_str,
                        context=context,
                        reward=reward,
                        round_num=round_num,
                    ),
                )
            else:
                return None

            # Convert result to ChallengerFeedback
            if isinstance(result, ChallengerFeedback):
                return result
            elif isinstance(result, dict):
                return ChallengerFeedback(
                    critique=result.get("critique", ""),
                    alternative_approach=result.get("alternative_approach", ""),
                    suggested_improvements=result.get(
                        "improvements", result.get("suggested_improvements", [])
                    ),
                    confidence=result.get("confidence", 0.5),
                )
            elif hasattr(result, "critique"):
                # DSPy prediction object
                improvements = []
                if hasattr(result, "improvements"):
                    imp = result.improvements
                    if isinstance(imp, list):
                        improvements = imp
                    elif isinstance(imp, str):
                        improvements = [imp]

                return ChallengerFeedback(
                    critique=getattr(result, "critique", ""),
                    alternative_approach=getattr(result, "alternative_approach", ""),
                    suggested_improvements=improvements,
                )

            return None

        except Exception:
            # Return None on challenger error - don't break the loop
            return None

    def _format_hint(self, feedback: ChallengerFeedback) -> str:
        """
        Format feedback as a hint string for the next solver call.

        Args:
            feedback: Feedback from challenger.

        Returns:
            Formatted hint string.
        """
        # Defensive type check for unexpected input
        if not isinstance(feedback, ChallengerFeedback):
            if feedback is None:
                return ""
            # If it's a string or other type, just convert to string
            return str(feedback)

        parts = []

        if feedback.critique:
            parts.append(f"Critique: {feedback.critique}")

        if feedback.alternative_approach:
            parts.append(f"Alternative approach: {feedback.alternative_approach}")

        if feedback.suggested_improvements:
            improvements = "; ".join(feedback.suggested_improvements)
            parts.append(f"Suggested improvements: {improvements}")

        return " | ".join(parts) if parts else ""
