"""
OWL Agent Implementation
This module implements the core functionalities for both user and assistant agents.
"""

from typing import Dict, List, Optional, Any
from pydantic import BaseModel
from camel.agents import BaseAgent
from camel.messages import BaseMessage
from owl.toolkits.base import BaseToolkit

class AgentRole(BaseModel):
    """Defines the role and capabilities of an agent."""
    name: str
    description: str
    capabilities: List[str]

class OwlRolePlaying:
    """Manages the interaction between user and assistant agents."""
    
    def __init__(
        self,
        user_role: AgentRole,
        assistant_role: AgentRole,
        toolkits: Optional[List[BaseToolkit]] = None
    ):
        """Initialize the role-playing session."""
        self.user_role = user_role
        self.assistant_role = assistant_role
        self.toolkits = toolkits or []
        
        # Initialize agents
        self.user_agent = self._create_user_agent()
        self.assistant_agent = self._create_assistant_agent()
        
    def _create_user_agent(self) -> BaseAgent:
        """Create and configure the user agent."""
        return BaseAgent.from_role(
            name=self.user_role.name,
            description=self.user_role.description,
            instructions=self._get_user_instructions()
        )
    
    def _create_assistant_agent(self) -> BaseAgent:
        """Create and configure the assistant agent."""
        return BaseAgent.from_role(
            name=self.assistant_role.name,
            description=self.assistant_role.description,
            instructions=self._get_assistant_instructions()
        )
    
    def _get_user_instructions(self) -> str:
        """Get the instructions for the user agent."""
        return f"""You are {self.user_role.name}. Your role is to:
- Clearly communicate task requirements
- Provide necessary context and information
- Review and validate assistant's responses
- Guide the conversation towards task completion"""
    
    def _get_assistant_instructions(self) -> str:
        """Get the instructions for the assistant agent."""
        toolkit_names = [t.name for t in self.toolkits]
        return f"""You are {self.assistant_role.name}. Your role is to:
- Understand and analyze user tasks
- Select appropriate tools from available toolkits: {', '.join(toolkit_names)}
- Execute tasks efficiently and accurately
- Provide clear explanations of actions and results"""
    
    def process_user_input(self, user_input: str) -> Dict[str, Any]:
        """Process user input and generate assistant response."""
        # Create user message
        user_message = BaseMessage(role="user", content=user_input)
        
        # Get assistant response
        assistant_response = self.assistant_agent.respond(
            user_message,
            toolkits=self.toolkits
        )
        
        return {
            "response": assistant_response.content,
            "used_tools": assistant_response.metadata.get("used_tools", []),
            "status": "success"
        }
    
    def get_available_tools(self) -> List[Dict[str, str]]:
        """Get list of available tools from all toolkits."""
        tools = []
        for toolkit in self.toolkits:
            tools.extend(toolkit.list_tools())
        return tools