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 ToolUsePattern(BasePattern):
    """
    Tool use pattern for agent interaction with external tools.

    This pattern enables agents to leverage external tools and APIs to complete
    tasks that require external information or capabilities. The agent can
    intelligently select and use appropriate tools based on the task requirements,
    parse tool results, and integrate them into the final response.

    The tool use pattern is essential for tasks requiring:
    - External data retrieval (APIs, databases, search engines)
    - Computations (calculators, data processors)
    - Actions (file operations, notifications, commands)

    The pattern orchestrates a conversation loop where the agent can:
    1. Analyze the task and available tools
    2. Select and invoke appropriate tools with parameters
    3. Process tool results and integrate them
    4. Iterate until the task is complete
    """

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

        :param agent: BaseAgent - The agent to use for tool selection and result processing.
        :param tools: List[BaseTool] - List of available tools the agent can use.
        :param config: Optional[PatternConfig] - Pattern configuration. If None, uses default with max_iterations=5.
        """
        if config is None:
            config = PatternConfig(
                name="tool_use", pattern_type=PatternType.TOOL_USE, max_iterations=5
            )
        super().__init__(config)
        self.agent = agent
        self.tools = {tool.name: tool for tool in tools}
        self._conversation = Conversation(id=f"tool-use-{id(self)}")
        self._state = State()

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

        The method enables the agent to iteratively select and use tools to complete
        the task. The agent receives descriptions of available tools and can invoke
        them by responding with specific commands. Tool results are integrated into
        the conversation until the agent provides a final answer.

        :param task: str - The task or question to execute with tool assistance.
        :param context: Optional[Dict[str, Any]] - Additional context for execution. Defaults to None.
        :return: PatternResult - Contains success status, final answer, tools used, and metadata.
        :raises ToolExecutionError: If a tool execution fails.

        Example:
            >>> tools = [CalculatorTool(), SearchTool()]
            >>> pattern = ToolUsePattern(agent=my_agent, tools=tools)
            >>> result = await pattern.execute("What is 15% of 240?")
            >>> print(f"Answer: {result.result}, Tools used: {result.metadata['tools_used']}")
        """
        self.reset()
        self._state.update_status(StateStatus.RUNNING)

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

        try:
            tools_description = self._get_tools_description()
            system_prompt = (
                f"You have access to the following tools:\n{tools_description}\n\n"
                "To use a tool, respond with: TOOL_CALL: <tool_name> <parameters>\n"
                "When you have the final answer, respond with: FINAL_ANSWER: <answer>"
            )

            self._conversation.add_message(Message(role=MessageRole.SYSTEM, content=system_prompt))
            self._conversation.add_message(Message(role=MessageRole.USER, content=task))

            final_answer = None

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

                state_data = {
                    "iteration": self.current_iteration,
                    "task": task,
                    "conversation": self._conversation.to_dict_list(),
                }

                step_result = await self.step(state_data)

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

            self._state.update_status(StateStatus.COMPLETED)

            return PatternResult(
                success=True,
                result=final_answer or "No final answer provided",
                iterations=self.current_iteration,
                metadata={
                    "tools_used": step_result.get("tools_used", []),
                    "conversation_length": len(self._conversation.messages),
                },
            )

        except Exception as e:
            logger.error(f"Tool use 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]:
        conversation = state.get("conversation", [])
        conversation_str = "\n".join([f"{m['role']}: {m['content']}" for m in conversation[-5:]])

        response = await self.agent.generate(conversation_str)

        if response.startswith("TOOL_CALL:"):
            tool_call = response[len("TOOL_CALL:") :].strip()
            parts = tool_call.split(maxsplit=1)
            tool_name = parts[0]
            tool_params = parts[1] if len(parts) > 1 else ""

            logger.debug(f"Calling tool: {tool_name} with params: {tool_params}")

            if tool_name in self.tools:
                tool_result = await self.tools[tool_name].execute(input=tool_params)
                result_msg = f"Tool '{tool_name}' result: {tool_result.result}"
                self._conversation.add_message(Message(role=MessageRole.TOOL, content=result_msg))

                return {
                    "is_final": False,
                    "answer": None,
                    "tool_called": tool_name,
                    "tools_used": [tool_name],
                }
            else:
                error_msg = f"Unknown tool: {tool_name}"
                self._conversation.add_message(Message(role=MessageRole.SYSTEM, content=error_msg))
                return {"is_final": False, "answer": None, "tools_used": []}

        elif response.startswith("FINAL_ANSWER:"):
            answer = response[len("FINAL_ANSWER:") :].strip()
            self._conversation.add_message(Message(role=MessageRole.ASSISTANT, content=answer))
            return {"is_final": True, "answer": answer, "tools_used": []}

        else:
            self._conversation.add_message(Message(role=MessageRole.ASSISTANT, content=response))
            return {"is_final": False, "answer": None, "tools_used": []}

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

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