"""
A2A Protocol Types

Pydantic models for the Agent2Agent (A2A) protocol.
Based on the A2A Protocol Specification v1.0.
"""

from enum import Enum
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, ConfigDict, Field


# =============================================================================
# Enums
# =============================================================================

class TaskState(str, Enum):
    """
    Defines the possible lifecycle states of a Task.
    
    Based on A2A Protocol Specification Section 4.1.3.
    """
    SUBMITTED = "submitted"
    WORKING = "working"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    INPUT_REQUIRED = "input_required"
    AUTH_REQUIRED = "auth_required"


class Role(str, Enum):
    """
    Defines the sender of a message in A2A protocol communication.
    
    Based on A2A Protocol Specification Section 4.1.5.
    """
    USER = "user"
    AGENT = "agent"


# =============================================================================
# Part Types
# =============================================================================

class TextPart(BaseModel):
    """
    A text content part.
    
    Based on A2A Protocol Specification Section 4.1.6.
    """
    text: str
    metadata: Optional[Dict[str, Any]] = None


class FilePart(BaseModel):
    """
    A file content part.
    
    Can contain either a URI reference or inline bytes.
    Based on A2A Protocol Specification Section 4.1.7.
    """
    file_uri: Optional[str] = Field(default=None, alias="fileWithUri")
    file_bytes: Optional[bytes] = Field(default=None, alias="fileWithBytes")
    media_type: Optional[str] = Field(default=None, alias="mediaType")
    name: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    
    model_config = ConfigDict(populate_by_name=True)


class DataPart(BaseModel):
    """
    A structured data part (JSON blob).
    
    Based on A2A Protocol Specification Section 4.1.8.
    """
    data: Dict[str, Any]
    metadata: Optional[Dict[str, Any]] = None


# Union type for all Part types
Part = Union[TextPart, FilePart, DataPart]


# =============================================================================
# Core Types
# =============================================================================

class Message(BaseModel):
    """
    A communication turn between client and server.
    
    Based on A2A Protocol Specification Section 4.1.4.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    message_id: str = Field(alias="messageId")
    role: Role
    parts: List[Part]
    context_id: Optional[str] = Field(default=None, alias="contextId")
    task_id: Optional[str] = Field(default=None, alias="taskId")
    metadata: Optional[Dict[str, Any]] = None
    extensions: Optional[Dict[str, str]] = None
    reference_task_ids: Optional[List[str]] = Field(default=None, alias="referenceTaskIds")


class TaskStatus(BaseModel):
    """
    Container for the status of a task.
    
    Based on A2A Protocol Specification Section 4.1.2.
    """
    state: TaskState
    message: Optional[Message] = None
    timestamp: Optional[str] = None


class Artifact(BaseModel):
    """
    An output generated by the agent as a result of a task.
    
    Based on A2A Protocol Specification Section 4.1.9.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    artifact_id: str = Field(alias="artifactId")
    name: Optional[str] = None
    description: Optional[str] = None
    parts: List[Part]
    metadata: Optional[Dict[str, Any]] = None
    extensions: Optional[Dict[str, str]] = None


class Task(BaseModel):
    """
    The core unit of action for A2A.
    
    Based on A2A Protocol Specification Section 4.1.1.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    id: str
    context_id: Optional[str] = Field(default=None, alias="contextId")
    status: TaskStatus
    artifacts: Optional[List[Artifact]] = None
    history: Optional[List[Message]] = None
    metadata: Optional[Dict[str, Any]] = None


# =============================================================================
# Agent Card Types
# =============================================================================

class AgentSkill(BaseModel):
    """
    Describes a skill/capability of an agent.
    
    Based on A2A Protocol Specification Section 4.4.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    id: str
    name: str
    description: str
    tags: Optional[List[str]] = None
    examples: Optional[List[str]] = None
    input_modes: Optional[List[str]] = Field(default=None, alias="inputModes")
    output_modes: Optional[List[str]] = Field(default=None, alias="outputModes")


class AgentCapabilities(BaseModel):
    """
    Describes the capabilities of an agent.
    
    Based on A2A Protocol Specification Section 4.4.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    streaming: bool = False
    push_notifications: bool = Field(default=False, alias="pushNotifications")
    state_transition_history: bool = Field(default=False, alias="stateTransitionHistory")


class AgentCard(BaseModel):
    """
    JSON metadata document describing an agent.
    
    Published by an A2A Server for discovery.
    Based on A2A Protocol Specification Section 8.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    name: str
    url: str
    version: str
    description: Optional[str] = None
    capabilities: AgentCapabilities
    skills: Optional[List[AgentSkill]] = None
    default_input_modes: Optional[List[str]] = Field(default=None, alias="defaultInputModes")
    default_output_modes: Optional[List[str]] = Field(default=None, alias="defaultOutputModes")
    provider: Optional[Dict[str, str]] = None
    documentation_url: Optional[str] = Field(default=None, alias="documentationUrl")


# =============================================================================
# Request/Response Types
# =============================================================================

class MessageConfig(BaseModel):
    """
    Configuration for message sending.
    
    Based on A2A Protocol Specification Section 3.2.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    accept_content_types: Optional[List[str]] = Field(default=None, alias="acceptContentTypes")
    blocking: Optional[bool] = None
    history_length: Optional[int] = Field(default=None, alias="historyLength")


class SendMessageRequest(BaseModel):
    """
    Request to send a message to an agent.
    
    Based on A2A Protocol Specification Section 3.1.1.
    """
    message: Message
    configuration: Optional[MessageConfig] = None
    metadata: Optional[Dict[str, Any]] = None


# =============================================================================
# Streaming Event Types
# =============================================================================

class TaskStatusUpdateEvent(BaseModel):
    """
    Event sent by agent to notify client of task status change.
    
    Based on A2A Protocol Specification Section 4.2.1.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    task_id: str = Field(alias="taskId")
    context_id: Optional[str] = Field(default=None, alias="contextId")
    status: TaskStatus
    final: bool = False
    metadata: Optional[Dict[str, Any]] = None


class TaskArtifactUpdateEvent(BaseModel):
    """
    Event sent when an artifact is generated.
    
    Based on A2A Protocol Specification Section 4.2.2.
    """
    model_config = ConfigDict(populate_by_name=True)
    
    task_id: str = Field(alias="taskId")
    context_id: Optional[str] = Field(default=None, alias="contextId")
    artifact: Artifact
    append: bool = False
    last_chunk: bool = Field(default=False, alias="lastChunk")
    metadata: Optional[Dict[str, Any]] = None
