"""
Tool Use Pattern - Data Processing

Demonstrates using tools for data validation, transformation, and processing tasks.
"""
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 DataValidatorTool(BaseTool):
    def __init__(self):
        config = ToolConfig(
            name="data_validator",
            description="Validates and cleans data entries",
            category=ToolCategory.COMPUTATION,
            parameters_schema={"data": "list[dict]", "rules": "dict"}
        )
        super().__init__(config)
    
    def validate_parameters(self, parameters: Dict[str, Any]) -> bool:
        return "data" in parameters and isinstance(parameters["data"], list)
    
    async def execute(self, data: list, rules: dict = None, **kwargs) -> ToolResult:
        try:
            valid_entries = []
            invalid_entries = []
            
            for entry in data:
                is_valid = True
                if rules and "required_fields" in rules:
                    for field in rules["required_fields"]:
                        if field not in entry or not entry[field]:
                            is_valid = False
                            break
                
                if is_valid:
                    valid_entries.append(entry)
                else:
                    invalid_entries.append(entry)
            
            return ToolResult(
                success=True,
                result={
                    "valid": valid_entries,
                    "invalid": invalid_entries,
                    "valid_count": len(valid_entries),
                    "invalid_count": len(invalid_entries)
                },
                metadata={"total_processed": len(data)}
            )
        except Exception as e:
            return ToolResult(success=False, result=None, error=str(e))


class DataTransformerTool(BaseTool):
    def __init__(self):
        config = ToolConfig(
            name="data_transformer",
            description="Transforms data into different formats",
            category=ToolCategory.COMPUTATION,
            parameters_schema={"data": "list", "transform_type": "string"}
        )
        super().__init__(config)
    
    def validate_parameters(self, parameters: Dict[str, Any]) -> bool:
        return "data" in parameters and isinstance(parameters["data"], list)
    
    async def execute(self, data: list, transform_type: str = "uppercase", **kwargs) -> ToolResult:
        try:
            if transform_type == "uppercase":
                result = [str(item).upper() for item in data]
            elif transform_type == "lowercase":
                result = [str(item).lower() for item in data]
            elif transform_type == "numeric":
                result = [float(item) if str(item).replace('.', '').isdigit() else 0 for item in data]
            else:
                return ToolResult(success=False, result=None, error=f"Unknown transform: {transform_type}")
            
            return ToolResult(
                success=True,
                result=result,
                metadata={"transform_type": transform_type, "count": len(result)}
            )
        except Exception as e:
            return ToolResult(success=False, result=None, error=str(e))


async def main():
    agent_config = AgentConfig(
        name="data_processor",
        model="openai/gpt-4o-mini",
        temperature=0.2,
        system_prompt="You are a data processing assistant that validates and transforms data using available tools."
    )
    
    agent = GEAIAgent(config=agent_config)
    
    pattern_config = PatternConfig(
        name="data_processing",
        pattern_type=PatternType.TOOL_USE,
        max_iterations=3
    )
    
    validator = DataValidatorTool()
    transformer = DataTransformerTool()
    pattern = ToolUsePattern(agent=agent, config=pattern_config, tools=[validator, transformer])
    
    task = """
    Process this customer data:
    - Validate entries have required fields (name, email)
    - Transform valid emails to lowercase
    
    Data: [
        {"name": "John Doe", "email": "JOHN@EXAMPLE.COM"},
        {"name": "Jane Smith", "email": ""},
        {"email": "bob@test.com"}
    ]
    """
    
    print(f"Data Processing Task:\n{task}\n")
    
    result = await pattern.execute(task)
    
    print(f"\nProcessing Complete: {result.success}")
    print(f"Result: {result.result}")


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