"""
Tool Use Pattern - Calculator Operations

Demonstrates using the ToolUsePattern to enable agents to use external tools
for mathematical calculations and data processing.
"""
import asyncio
from typing import Dict, Any
from pygeai_orchestration.core.base import (
    AgentConfig, PatternConfig, PatternType,
    BaseTool, ToolConfig, ToolResult, ToolCategory
)
from pygeai_orchestration.core.base.geai_agent import GEAIAgent
from pygeai_orchestration.patterns import ToolUsePattern


class CalculatorTool(BaseTool):
    def __init__(self):
        config = ToolConfig(
            name="calculator",
            description="Performs mathematical calculations",
            category=ToolCategory.COMPUTATION,
            parameters_schema={"operation": "string", "values": "list[float]"}
        )
        super().__init__(config)
    
    def validate_parameters(self, parameters: Dict[str, Any]) -> bool:
        return "operation" in parameters and "values" in parameters
    
    async def execute(self, operation: str, values: list, **kwargs) -> ToolResult:
        try:
            if operation == "add":
                result = sum(values)
            elif operation == "multiply":
                result = 1
                for v in values:
                    result *= v
            elif operation == "average":
                result = sum(values) / len(values) if values else 0
            elif operation == "max":
                result = max(values) if values else 0
            elif operation == "min":
                result = min(values) if values else 0
            else:
                return ToolResult(
                    success=False,
                    result=None,
                    error=f"Unknown operation: {operation}"
                )
            
            return ToolResult(
                success=True,
                result=result,
                metadata={"operation": operation, "input_count": len(values)}
            )
        except Exception as e:
            return ToolResult(success=False, result=None, error=str(e))


async def main():
    agent_config = AgentConfig(
        name="calculator_agent",
        model="openai/gpt-4o-mini",
        temperature=0.1,
        system_prompt="You are a mathematical assistant that uses tools to perform calculations."
    )
    
    agent = GEAIAgent(config=agent_config)
    
    pattern_config = PatternConfig(
        name="calculator_usage",
        pattern_type=PatternType.TOOL_USE,
        max_iterations=3
    )
    
    calculator = CalculatorTool()
    pattern = ToolUsePattern(agent=agent, config=pattern_config, tools=[calculator])
    
    task = "Calculate the average of these sales figures: 1250, 3400, 2100, 4750, 1890"
    print(f"Task: {task}\n")
    
    result = await pattern.execute(task)
    
    print(f"\nSuccess: {result.success}")
    print(f"Result: {result.result}")


if __name__ == "__main__":
    asyncio.run(main())
