"""
Agent wrapper for simplified agent creation and management.
"""

import uuid
from typing import Any, Callable, Coroutine

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.types import AgentCard, AgentSkill, Message, Task
from a2a.utils.message import new_agent_text_message

from synqed.agent_card import AgentCardBuilder


class Agent:
    """
    Simplified agent wrapper that handles A2A agent creation and execution.
    
    This class provides a high-level interface for creating agents without
    needing to understand the underlying A2A protocol details.
    
    Example:
        ```python
        async def handle_request(context):
            # Your agent logic here
            return "Response"
        
        agent = Agent(
            name="My Agent",
            description="Does something useful",
            skills=["skill1", "skill2"],
            executor=handle_request
        )
        ```
    """
    
    def __init__(
        self,
        name: str,
        description: str,
        skills: list[str] | list[dict[str, Any]],
        executor: Callable[[RequestContext], Coroutine[Any, Any, str | Message]] | None = None,
        version: str = "1.0.0",
        url: str | None = None,
        default_input_modes: list[str] | None = None,
        default_output_modes: list[str] | None = None,
        capabilities: dict[str, Any] | None = None,
        security_schemes: dict[str, Any] | None = None,
    ):
        """
        Initialize an Agent.
        
        Args:
            name: Human-readable name for the agent
            description: Description of what the agent does
            skills: List of skill names (strings) or skill dictionaries with details
            executor: Optional async function to handle requests
            version: Agent version (default: "1.0.0")
            url: Base URL for the agent (set automatically if using AgentServer)
            default_input_modes: Supported input MIME types
            default_output_modes: Supported output MIME types
            capabilities: Agent capabilities (streaming, push_notifications, etc.)
            security_schemes: Security configuration
        """
        self.name = name
        self.description = description
        self.version = version
        self.url = url or f"http://localhost:8000"  # Default, can be overridden
        self._skills_config = skills
        self._executor_func = executor
        self._agent_executor: AgentExecutor | None = None
        
        # Build agent card
        self._card_builder = AgentCardBuilder(
            name=name,
            description=description,
            version=version,
            url=self.url,
        )
        
        # Add skills
        for skill in skills:
            if isinstance(skill, str):
                self._card_builder.add_skill(
                    skill_id=skill,
                    name=skill.replace("_", " ").title(),
                    description=f"Skill for {skill}",
                    tags=[skill]
                )
            else:
                self._card_builder.add_skill(**skill)
        
        # Set optional configurations
        if default_input_modes:
            self._card_builder.set_default_input_modes(default_input_modes)
        if default_output_modes:
            self._card_builder.set_default_output_modes(default_output_modes)
        if capabilities:
            self._card_builder.set_capabilities(**capabilities)
        if security_schemes:
            for name, scheme in security_schemes.items():
                self._card_builder.add_security_scheme(name, scheme)
    
    @property
    def card(self) -> AgentCard:
        """Get the agent's card."""
        return self._card_builder.build()
    
    @property
    def skills(self) -> list[str] | list[dict[str, Any]]:
        """Get the agent's skills configuration."""
        return self._skills_config
    
    def set_url(self, url: str) -> None:
        """Update the agent's URL."""
        self.url = url
        self._card_builder.set_url(url)
    
    def get_executor(self) -> AgentExecutor:
        """
        Get or create the agent executor.
        
        Returns:
            AgentExecutor instance for handling requests
        """
        if self._agent_executor is None:
            if self._executor_func is None:
                # Create a default executor
                self._agent_executor = self._create_default_executor()
            else:
                # Wrap the user's function in an AgentExecutor
                self._agent_executor = self._create_custom_executor(self._executor_func)
        
        return self._agent_executor
    
    def set_executor(
        self, 
        executor: Callable[[RequestContext], Coroutine[Any, Any, str | Message]]
    ) -> None:
        """
        Set a custom executor function for this agent.
        
        Args:
            executor: Async function that takes RequestContext and returns a response
        """
        self._executor_func = executor
        self._agent_executor = self._create_custom_executor(executor)
    
    def _create_default_executor(self) -> AgentExecutor:
        """Create a default executor that returns a simple response."""
        from a2a.server.events import EventQueue
        from a2a.types import TaskStatus, TaskState, TaskStatusUpdateEvent
        
        agent = self  # Capture self for closure
        
        class DefaultExecutor(AgentExecutor):
            async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
                # Extract the user's message
                user_message = context.get_user_input()
                
                # Create a simple response
                response_text = (
                    f"Hello! I am {agent.name}. "
                    f"I received your message: '{user_message}'. "
                    f"I can help with: {', '.join([s if isinstance(s, str) else s.get('id', 'unknown') for s in agent._skills_config])}"
                )
                
                message = new_agent_text_message(
                    text=response_text,
                    task_id=context.task_id
                )
                
                # Publish the message to the event queue
                await event_queue.enqueue_event(message)
                
                # Mark the task as completed
                task_status = TaskStatus(state=TaskState.completed)
                event = TaskStatusUpdateEvent(
                    task_id=context.task_id,
                    context_id=context.context_id,
                    status=task_status,
                    final=True
                )
                await event_queue.enqueue_event(event)
            
            async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
                # Simple cancellation - just mark as canceled
                from a2a.types import TaskStatus, TaskState, TaskStatusUpdateEvent
                
                task_status = TaskStatus(state=TaskState.canceled)
                event = TaskStatusUpdateEvent(
                    task_id=context.task_id,
                    context_id=context.context_id,
                    status=task_status,
                    final=True
                )
                await event_queue.enqueue_event(event)
        
        return DefaultExecutor()
    
    def _create_custom_executor(
        self,
        handler: Callable[[RequestContext], Coroutine[Any, Any, str | Message]]
    ) -> AgentExecutor:
        """Create an executor from a user-provided handler function."""
        from a2a.server.events import EventQueue
        from a2a.types import TaskStatus, TaskState, TaskStatusUpdateEvent
        
        class CustomExecutor(AgentExecutor):
            async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
                result = await handler(context)
                
                # If the handler returns a string, wrap it in a Message
                if isinstance(result, str):
                    message = new_agent_text_message(
                        text=result,
                        task_id=context.task_id
                    )
                else:
                    message = result
                
                # Publish the message to the event queue
                await event_queue.enqueue_event(message)
                
                # Mark the task as completed
                task_status = TaskStatus(state=TaskState.completed)
                event = TaskStatusUpdateEvent(
                    task_id=context.task_id,
                    context_id=context.context_id,
                    status=task_status,
                    final=True
                )
                await event_queue.enqueue_event(event)
            
            async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
                # Simple cancellation - just mark as canceled
                from a2a.types import TaskStatus, TaskState, TaskStatusUpdateEvent
                
                task_status = TaskStatus(state=TaskState.canceled)
                event = TaskStatusUpdateEvent(
                    task_id=context.task_id,
                    context_id=context.context_id,
                    status=task_status,
                    final=True
                )
                await event_queue.enqueue_event(event)
        
        return CustomExecutor()
    
    def __repr__(self) -> str:
        """String representation of the agent."""
        return (
            f"Agent(name='{self.name}', "
            f"skills={len(self._skills_config)}, "
            f"url='{self.url}')"
        )

