"""AI agent system implementation"""

from typing import List, Dict, Any, Optional, Generator
from diona.ai.client.simplecli.models import Message
from diona.ai.client.simplecli.agent.prompt import AgentPromptBuilder
from diona.ai.client.simplecli.agent.parser import ToolCallParser


class AIAgentSystem:
    """AI agent system for coordinating AI and tool interactions"""
    
    def __init__(self, provider, tool_manager, config: Dict[str, Any]):
        """Initialize AI agent system
        
        Args:
            provider: AI provider (ProviderManager)
            tool_manager: Tool manager
            config: Agent configuration
        """
        self.provider = provider
        self.tool_manager = tool_manager
        self.config = config
        self.prompt_builder = AgentPromptBuilder(config.get('system-prompt', ''))
        self.parser = ToolCallParser()
        self.max_tool_calls = config.get('max-tool-calls', 5)
    
    def run_agent(self, messages: List[Message]) -> Generator[str, None, None]:
        """Run agent to process user request
        
        Args:
            messages: List of messages
            
        Yields:
            Response chunks
        """
        # Get available tools
        tools = self._get_available_tools()
        
        # Build system prompt with tool information
        system_prompt = self.prompt_builder.build_prompt(tools)
        
        # Create messages with system prompt
        agent_messages = [Message("system", system_prompt)] + messages
        
        # Process with tool calls
        tool_calls = 0
        while tool_calls < self.max_tool_calls:
            # Get AI response
            response = ""
            for chunk in self.provider.chat_stream(agent_messages, self.provider.model, self.config.get('temperature', 0.7)):
                response += chunk
                yield chunk
            
            # Check if it's a tool call
            tool_call = self.parser.parse(response)
            if not tool_call:
                # No tool call, return the response
                break
            
            # Execute tool
            tool_name = tool_call.get('tool')
            params = tool_call.get('params', {})
            
            # Execute tool
            tool_result = self.tool_manager.execute_tool(tool_name, **params)
            
            # Format tool response
            tool_response = self.parser.format_tool_response(tool_name, tool_result)
            
            # Yield the tool result to display to the user
            yield f"\nTool Result: {tool_result}\n"
            
            # Add tool response to messages
            agent_messages.append(Message("assistant", response))
            agent_messages.append(Message("user", tool_response))
            
            tool_calls += 1
    
    def handle_tool_calls(self, tool_calls: Dict[str, Any]) -> str:
        """Handle tool calls from AI
        
        Args:
            tool_calls: Tool call data
            
        Returns:
            Tool execution result
        """
        tool_name = tool_calls.get('tool')
        params = tool_calls.get('params', {})
        
        return self.tool_manager.execute_tool(tool_name, **params)
    
    def format_tool_response(self, tool_result: Any) -> str:
        """Format tool response for AI
        
        Args:
            tool_result: Tool execution result
            
        Returns:
            Formatted tool response
        """
        return str(tool_result)
    
    def generate_agent_prompt(self) -> str:
        """Generate agent system prompt
        
        Returns:
            Agent system prompt
        """
        tools = self._get_available_tools()
        return self.prompt_builder.build_prompt(tools)
    
    def _get_available_tools(self) -> List[Dict[str, Any]]:
        """Get available tools with schemas
        
        Returns:
            List of tool schemas
        """
        tools = []
        for tool_name in self.tool_manager.get_available_tools():
            schema = self.tool_manager.get_tool_schema(tool_name)
            if schema:
                tools.append(schema)
        return tools
