from typing import Any, Dict, Optional
from pygeai_orchestration.core.base import BasePattern, PatternConfig, PatternResult, PatternType
from pygeai_orchestration.core.common import Message, MessageRole, Conversation, State, StateStatus
import logging

logger = logging.getLogger("pygeai_orchestration")


class ReflectionPattern(BasePattern):
    """
    Reflection pattern for iterative self-improvement of agent responses.

    This pattern enables agents to critique and refine their own outputs through
    multiple iterations. Each iteration consists of generating a response, reflecting
    on its quality, and producing an improved version. The process continues until
    convergence or max_iterations is reached.

    The reflection pattern is particularly effective for tasks requiring high-quality
    outputs such as content creation, analysis, and problem-solving where iterative
    refinement leads to better results.
    """

    def __init__(
        self, agent, config: Optional[PatternConfig] = None, reflection_prompt: Optional[str] = None
    ):
        """
        Initialize the Reflection pattern.

        :param agent: BaseAgent - The agent to use for generation and reflection.
        :param config: Optional[PatternConfig] - Pattern configuration. If None, uses default with max_iterations=3.
        :param reflection_prompt: Optional[str] - Custom reflection prompt. If None, uses default prompt.
        """
        if config is None:
            config = PatternConfig(
                name="reflection", pattern_type=PatternType.REFLECTION, max_iterations=3
            )
        super().__init__(config)
        self.agent = agent
        self.reflection_prompt = reflection_prompt or (
            "Review your previous response. Identify any issues, inaccuracies, or areas for improvement. "
            "Provide a refined and improved version of your response."
        )
        self._conversation = Conversation(id=f"reflection-{id(self)}")
        self._state = State()

    async def execute(self, task: str, context: Optional[Dict[str, Any]] = None) -> PatternResult:
        """
        Execute the reflection pattern on the given task.

        The method iteratively improves the agent's response through self-reflection.
        Each iteration generates a response, reflects on it, and produces an improved
        version. The process stops when max_iterations is reached or convergence is detected.

        :param task: str - The task or prompt to execute.
        :param context: Optional[Dict[str, Any]] - Additional context for execution. Defaults to None.
        :return: PatternResult - Contains success status, final refined result, iteration count, and metadata.
        :raises ValueError: If task is empty or None.

        Example:
            >>> pattern = ReflectionPattern(agent=my_agent)
            >>> result = await pattern.execute("Explain quantum computing in simple terms")
            >>> print(f"Result after {result.iterations} iterations: {result.result}")
        """
        self.reset()
        self._state.update_status(StateStatus.RUNNING)

        logger.info(f"Starting reflection pattern for task: {task[:50]}...")

        try:
            self._conversation.add_message(Message(role=MessageRole.USER, content=task))

            current_response = None

            while self.current_iteration < self.config.max_iterations:
                self.increment_iteration()
                logger.debug(
                    f"Reflection iteration {self.current_iteration}/{self.config.max_iterations}"
                )

                state_data = {
                    "iteration": self.current_iteration,
                    "task": task,
                    "current_response": current_response,
                }

                step_result = await self.step(state_data)
                current_response = step_result.get("response")

                if step_result.get("should_stop", False):
                    logger.info(f"Reflection converged at iteration {self.current_iteration}")
                    break

            self._state.update_status(StateStatus.COMPLETED)

            return PatternResult(
                success=True,
                result=current_response,
                iterations=self.current_iteration,
                metadata={
                    "conversation_length": len(self._conversation.messages),
                    "final_iteration": self.current_iteration,
                },
            )

        except Exception as e:
            logger.error(f"Reflection pattern failed: {str(e)}")
            self._state.update_status(StateStatus.FAILED)
            return PatternResult(
                success=False, result=None, iterations=self.current_iteration, error=str(e)
            )

    async def step(self, state: Dict[str, Any]) -> Dict[str, Any]:
        iteration = state.get("iteration", 1)
        task = state.get("task")
        current_response = state.get("current_response")

        if iteration == 1:
            logger.debug("Generating initial response")
            response = await self.agent.generate(task)
        else:
            logger.debug(f"Reflecting on iteration {iteration}")
            reflection_request = (
                f"Previous response:\n{current_response}\n\n{self.reflection_prompt}"
            )
            response = await self.agent.generate(reflection_request)

        self._conversation.add_message(Message(role=MessageRole.ASSISTANT, content=response))

        should_stop = iteration >= self.config.max_iterations

        return {"response": response, "should_stop": should_stop, "iteration": iteration}

    def get_conversation(self) -> Conversation:
        return self._conversation
