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

logger = logging.getLogger("pygeai_orchestration")


class ReActPattern(BasePattern):
    """
    ReAct (Reasoning and Acting) pattern for structured agent decision-making.

    This pattern implements the ReAct framework that interleaves reasoning traces
    and task-specific actions in a structured loop. The agent follows a think-act-observe
    cycle to solve complex tasks requiring iterative reasoning and tool use.

    The ReAct pattern is particularly effective for:
    - Multi-step reasoning tasks requiring careful thought
    - Tasks requiring dynamic tool selection based on intermediate results
    - Debugging and transparency (reasoning traces provide explainability)
    - Complex problem-solving requiring iterative refinement

    The pattern follows this cycle:
    1. **Thought**: Agent reasons about the current state and next steps
    2. **Action**: Agent selects and executes a tool or provides final answer
    3. **Observation**: Results are integrated and inform next thought

    Reference: Yao et al. (2022) "ReAct: Synergizing Reasoning and Acting in Language Models"
    """

    def __init__(
        self, agent, tools: Optional[List[BaseTool]] = None, config: Optional[PatternConfig] = None
    ):
        """
        Initialize the ReAct pattern.

        :param agent: BaseAgent - The agent to use for reasoning and action generation.
        :param tools: Optional[List[BaseTool]] - List of tools available to the agent. Defaults to None.
        :param config: Optional[PatternConfig] - Pattern configuration. If None, uses default with max_iterations=10.
        """
        if config is None:
            config = PatternConfig(name="react", pattern_type=PatternType.REACT, max_iterations=10)
        super().__init__(config)
        self.agent = agent
        self.tools = {tool.name: tool for tool in (tools or [])}
        self._conversation = Conversation(id=f"react-{id(self)}")
        self._state = State()
        self._thought_history: List[str] = []
        self._action_history: List[str] = []
        self._observation_history: List[str] = []

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

        The method orchestrates the think-act-observe cycle where the agent iteratively
        reasons about the task, selects actions (tool calls), observes results, and
        refines its approach until reaching a final answer.

        :param task: str - The task or question to solve using ReAct reasoning.
        :param context: Optional[Dict[str, Any]] - Additional context for execution. Defaults to None.
        :return: PatternResult - Contains success status, final answer, and complete reasoning trace.
        :raises PatternExecutionError: If the ReAct pattern execution fails.

        Example:
            >>> tools = [SearchTool(), CalculatorTool()]
            >>> pattern = ReActPattern(agent=my_agent, tools=tools)
            >>> result = await pattern.execute("What's the population of Tokyo times 2?")
            >>> print(f"Answer: {result.result}")
            >>> print(f"Reasoning trace: {result.metadata['thoughts']}")
        """
        self.reset()
        self._thought_history.clear()
        self._action_history.clear()
        self._observation_history.clear()
        self._state.update_status(StateStatus.RUNNING)

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

        try:
            system_prompt = self._build_system_prompt()
            self._conversation.add_message(Message(role=MessageRole.SYSTEM, content=system_prompt))
            self._conversation.add_message(Message(role=MessageRole.USER, content=f"Task: {task}"))

            final_answer = None

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

                state_data = {
                    "iteration": self.current_iteration,
                    "task": task,
                    "history": {
                        "thoughts": self._thought_history,
                        "actions": self._action_history,
                        "observations": self._observation_history,
                    },
                }

                step_result = await self.step(state_data)

                if step_result.get("is_final"):
                    final_answer = step_result.get("answer")
                    logger.info(f"ReAct completed at iteration {self.current_iteration}")
                    break

            self._state.update_status(StateStatus.COMPLETED)

            return PatternResult(
                success=True,
                result=final_answer or "No final answer reached",
                iterations=self.current_iteration,
                metadata={
                    "thoughts": self._thought_history,
                    "actions": self._action_history,
                    "observations": self._observation_history,
                },
            )

        except Exception as e:
            logger.error(f"ReAct 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]:
        history = state.get("history", {})

        prompt = self._build_step_prompt(history)
        response = await self.agent.generate(prompt)

        thought, action, observation, is_final, answer = self._parse_response(response)

        if thought:
            self._thought_history.append(thought)
            logger.debug(f"Thought: {thought[:100]}...")

        if action:
            self._action_history.append(action)
            logger.debug(f"Action: {action[:100]}...")

            if not is_final:
                obs = await self._execute_action(action)
                self._observation_history.append(obs)
                logger.debug(f"Observation: {obs[:100]}...")

        return {
            "is_final": is_final,
            "answer": answer,
            "thought": thought,
            "action": action,
            "observation": observation,
        }

    def _build_system_prompt(self) -> str:
        tools_desc = self._get_tools_description() if self.tools else "No tools available."
        return (
            "You are a reasoning agent following the ReAct pattern.\n"
            "For each step, provide your response in this format:\n\n"
            "Thought: <your reasoning about what to do next>\n"
            "Action: <the action to take or ANSWER if you have the final answer>\n"
            "Observation: <what you learned from the action>\n\n"
            f"Available tools:\n{tools_desc}\n\n"
            "When you have the final answer, use: Action: ANSWER <your answer>"
        )

    def _build_step_prompt(self, history: Dict[str, Any]) -> str:
        thoughts = history.get("thoughts", [])
        actions = history.get("actions", [])
        observations = history.get("observations", [])

        history_str = ""
        for i, (t, a, o) in enumerate(zip(thoughts, actions, observations)):
            history_str += f"\nStep {i + 1}:\nThought: {t}\nAction: {a}\nObservation: {o}\n"

        return f"Previous steps:{history_str}\n\nProvide the next step:"

    def _parse_response(self, response: str) -> tuple:
        thought = ""
        action = ""
        observation = ""
        is_final = False
        answer = None

        lines = response.strip().split("\n")
        for line in lines:
            if line.startswith("Thought:"):
                thought = line[len("Thought:") :].strip()
            elif line.startswith("Action:"):
                action = line[len("Action:") :].strip()
                if action.startswith("ANSWER"):
                    is_final = True
                    answer = action[len("ANSWER") :].strip()
            elif line.startswith("Observation:"):
                observation = line[len("Observation:") :].strip()

        return thought, action, observation, is_final, answer

    async def _execute_action(self, action: str) -> str:
        parts = action.split(maxsplit=1)
        tool_name = parts[0]
        tool_input = parts[1] if len(parts) > 1 else ""

        if tool_name in self.tools:
            result = await self.tools[tool_name].execute(input=tool_input)
            return str(result.result) if result.success else f"Error: {result.error}"
        else:
            return f"Tool '{tool_name}' not found"

    def _get_tools_description(self) -> str:
        if not self.tools:
            return "None"
        descriptions = [f"- {tool.name}: {tool.description}" for tool in self.tools.values()]
        return "\n".join(descriptions)

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