"""
A2A Protocol Pydantic Models

Complete type definitions for the A2A Protocol (v0.3).
All models use Pydantic v2 for validation and serialization.

Based on: a2a_protocol_deep_dive.md and official a2a-python SDK types.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Annotated, Any, Literal, Union

from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator


# ═══════════════════════════════════════════════════════════════════════════════
# ENUMS
# ═══════════════════════════════════════════════════════════════════════════════


class TaskState(str, Enum):
    """All valid A2A task states."""

    SUBMITTED = "submitted"
    WORKING = "working"
    INPUT_REQUIRED = "input-required"
    AUTH_REQUIRED = "auth-required"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELED = "canceled"
    REJECTED = "rejected"


class MessageRole(str, Enum):
    """Role of message sender."""

    USER = "user"
    AGENT = "agent"


class PartKind(str, Enum):
    """Discriminator for Part types."""

    TEXT = "text"
    FILE = "file"
    DATA = "data"


class Transport(str, Enum):
    """Supported transport protocols."""

    JSONRPC = "JSONRPC"
    GRPC = "GRPC"
    HTTP_JSON = "HTTP+JSON"


class SecuritySchemeType(str, Enum):
    """OpenAPI 3.0 security scheme types."""

    API_KEY = "apiKey"
    HTTP = "http"
    OAUTH2 = "oauth2"
    OPENID_CONNECT = "openIdConnect"
    MUTUAL_TLS = "mutualTLS"


class ApiKeyLocation(str, Enum):
    """Where to send API key."""

    HEADER = "header"
    QUERY = "query"
    COOKIE = "cookie"


# ═══════════════════════════════════════════════════════════════════════════════
# BASE CONFIG
# ═══════════════════════════════════════════════════════════════════════════════


class A2AModel(BaseModel):
    """Base model with common configuration."""

    model_config = ConfigDict(
        extra="allow",  # Allow extra fields for forward compatibility
        populate_by_name=True,
        str_strip_whitespace=True,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# FILE TYPES
# ═══════════════════════════════════════════════════════════════════════════════


class FileWithUri(A2AModel):
    """File reference via URI."""

    uri: str
    name: str | None = None
    mime_type: str | None = Field(default=None, alias="mimeType")


class FileWithBytes(A2AModel):
    """File with inline base64 content."""

    bytes: str  # Base64 encoded
    name: str | None = None
    mime_type: str | None = Field(default=None, alias="mimeType")


FileContent = Union[FileWithUri, FileWithBytes]


# ═══════════════════════════════════════════════════════════════════════════════
# PART TYPES (Discriminated Union)
# ═══════════════════════════════════════════════════════════════════════════════


class TextPart(A2AModel):
    """Text content part."""

    kind: Literal["text"] = "text"
    text: str
    metadata: dict[str, Any] | None = None


class FilePart(A2AModel):
    """File content part."""

    kind: Literal["file"] = "file"
    file: FileWithUri | FileWithBytes
    metadata: dict[str, Any] | None = None


class DataPart(A2AModel):
    """Structured data part."""

    kind: Literal["data"] = "data"
    data: dict[str, Any]
    metadata: dict[str, Any] | None = None


Part = Annotated[
    Union[TextPart, FilePart, DataPart],
    Field(discriminator="kind"),
]


# ═══════════════════════════════════════════════════════════════════════════════
# MESSAGE & ARTIFACT
# ═══════════════════════════════════════════════════════════════════════════════


class Message(A2AModel):
    """A2A Message object - fundamental communication unit."""

    role: MessageRole
    message_id: str = Field(alias="messageId")
    parts: list[Part]
    context_id: str | None = Field(default=None, alias="contextId")
    task_id: str | None = Field(default=None, alias="taskId")
    reference_task_ids: list[str] | None = Field(default=None, alias="referenceTaskIds")
    extensions: list[str] | None = None
    metadata: dict[str, Any] | None = None
    kind: Literal["message"] = "message"


class Artifact(A2AModel):
    """Tangible output generated by a task."""

    artifact_id: str = Field(alias="artifactId")
    name: str | None = None
    description: str | None = None
    parts: list[Part]
    extensions: list[str] | None = None
    metadata: dict[str, Any] | None = None
    # Streaming fields
    append: bool | None = None
    last_chunk: bool | None = Field(default=None, alias="lastChunk")
    index: int | None = None


# ═══════════════════════════════════════════════════════════════════════════════
# TASK STATUS & TASK
# ═══════════════════════════════════════════════════════════════════════════════


class TaskStatus(A2AModel):
    """Current status of a task."""

    state: TaskState
    message: Message | None = None
    timestamp: datetime | str | None = None


class Task(A2AModel):
    """A2A Task - fundamental unit of work."""

    id: str
    context_id: str = Field(alias="contextId")
    status: TaskStatus
    artifacts: list[Artifact] | None = None
    history: list[Message] | None = None
    metadata: dict[str, Any] | None = None
    kind: Literal["task"] = "task"


# ═══════════════════════════════════════════════════════════════════════════════
# STREAMING EVENTS
# ═══════════════════════════════════════════════════════════════════════════════


class TaskStatusUpdateEvent(A2AModel):
    """SSE event for task status updates."""

    kind: Literal["status-update"] = "status-update"
    task_id: str = Field(alias="taskId")
    context_id: str = Field(alias="contextId")
    status: TaskStatus
    final: bool = False
    metadata: dict[str, Any] | None = None


class TaskArtifactUpdateEvent(A2AModel):
    """SSE event for artifact updates."""

    kind: Literal["artifact-update"] = "artifact-update"
    task_id: str = Field(alias="taskId")
    context_id: str = Field(alias="contextId")
    artifact: Artifact
    append: bool = False
    last_chunk: bool = Field(default=False, alias="lastChunk")
    metadata: dict[str, Any] | None = None


# ═══════════════════════════════════════════════════════════════════════════════
# SECURITY SCHEMES (OpenAPI 3.0)
# ═══════════════════════════════════════════════════════════════════════════════


class ApiKeySecurityScheme(A2AModel):
    """API Key authentication."""

    type: Literal["apiKey"] = "apiKey"
    name: str
    in_: ApiKeyLocation = Field(alias="in")
    description: str | None = None


class HttpSecurityScheme(A2AModel):
    """HTTP authentication (Bearer, Basic, etc)."""

    type: Literal["http"] = "http"
    scheme: str  # e.g., "Bearer", "Basic"
    bearer_format: str | None = Field(default=None, alias="bearerFormat")
    description: str | None = None


class OAuth2Flow(A2AModel):
    """OAuth2 flow configuration."""

    authorization_url: str | None = Field(default=None, alias="authorizationUrl")
    token_url: str | None = Field(default=None, alias="tokenUrl")
    refresh_url: str | None = Field(default=None, alias="refreshUrl")
    scopes: dict[str, str] = Field(default_factory=dict)


class OAuth2Flows(A2AModel):
    """OAuth2 flows container."""

    authorization_code: OAuth2Flow | None = Field(default=None, alias="authorizationCode")
    implicit: OAuth2Flow | None = None
    password: OAuth2Flow | None = None
    client_credentials: OAuth2Flow | None = Field(default=None, alias="clientCredentials")


class OAuth2SecurityScheme(A2AModel):
    """OAuth 2.0 authentication."""

    type: Literal["oauth2"] = "oauth2"
    flows: OAuth2Flows
    oauth2_metadata_url: str | None = Field(default=None, alias="oauth2MetadataUrl")
    description: str | None = None


class OpenIdConnectSecurityScheme(A2AModel):
    """OpenID Connect authentication."""

    type: Literal["openIdConnect"] = "openIdConnect"
    open_id_connect_url: str = Field(alias="openIdConnectUrl")
    description: str | None = None


class MutualTLSSecurityScheme(A2AModel):
    """Mutual TLS authentication."""

    type: Literal["mutualTLS"] = "mutualTLS"
    description: str | None = None


SecurityScheme = Union[
    ApiKeySecurityScheme,
    HttpSecurityScheme,
    OAuth2SecurityScheme,
    OpenIdConnectSecurityScheme,
    MutualTLSSecurityScheme,
]


# ═══════════════════════════════════════════════════════════════════════════════
# AGENT CARD
# ═══════════════════════════════════════════════════════════════════════════════


class AgentProvider(A2AModel):
    """Information about the agent's provider."""

    organization: str | None = None
    url: str | None = None


class AgentExtension(A2AModel):
    """Protocol extension declaration."""

    uri: str
    description: str | None = None
    required: bool = False


class AgentCapabilities(A2AModel):
    """Agent capabilities declaration."""

    streaming: bool = False
    push_notifications: bool = Field(default=False, alias="pushNotifications")
    state_transition_history: bool = Field(default=False, alias="stateTransitionHistory")
    extensions: list[AgentExtension] | None = None


class AgentSkill(A2AModel):
    """A specific skill/capability the agent provides."""

    id: str
    name: str
    description: str | None = None
    tags: list[str] = Field(default_factory=list)
    examples: list[str] | None = None
    input_modes: list[str] | None = Field(default=None, alias="inputModes")
    output_modes: list[str] | None = Field(default=None, alias="outputModes")
    security: list[dict[str, list[str]]] | None = None


class AgentInterface(A2AModel):
    """Additional protocol interface."""

    url: str
    transport: Transport
    description: str | None = None


class AgentCardSignature(A2AModel):
    """Digital signature for agent card verification."""

    algorithm: str
    key_id: str = Field(alias="keyId")
    signature: str


class AgentCard(A2AModel):
    """
    A2A Agent Card - the primary discovery document.

    Located at: /.well-known/agent-card.json
    """

    # Required fields
    name: str
    url: str
    version: str
    capabilities: AgentCapabilities
    skills: list[AgentSkill]

    # Recommended fields
    description: str | None = None
    protocol_version: str | None = Field(default=None, alias="protocolVersion")
    default_input_modes: list[str] | None = Field(default=None, alias="defaultInputModes")
    default_output_modes: list[str] | None = Field(default=None, alias="defaultOutputModes")
    provider: AgentProvider | None = None
    documentation_url: str | None = Field(default=None, alias="documentationUrl")

    # Optional fields
    preferred_transport: Transport | None = Field(default=None, alias="preferredTransport")
    additional_interfaces: list[AgentInterface] | None = Field(
        default=None, alias="additionalInterfaces"
    )
    icon_url: str | None = Field(default=None, alias="iconUrl")
    supports_authenticated_extended_card: bool | None = Field(
        default=None, alias="supportsAuthenticatedExtendedCard"
    )
    security_schemes: dict[str, SecurityScheme] | None = Field(
        default=None, alias="securitySchemes"
    )
    security: list[dict[str, list[str]]] | None = None
    signatures: list[AgentCardSignature] | None = None


# ═══════════════════════════════════════════════════════════════════════════════
# PUSH NOTIFICATIONS
# ═══════════════════════════════════════════════════════════════════════════════


class PushNotificationAuthentication(A2AModel):
    """Authentication config for push notifications."""

    schemes: list[str]
    credentials: str | None = None


class PushNotificationConfig(A2AModel):
    """Webhook configuration for push notifications."""

    id: str
    url: str
    token: str | None = None
    authentication: PushNotificationAuthentication | None = None


# ═══════════════════════════════════════════════════════════════════════════════
# MESSAGE CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════


class MessageSendConfiguration(A2AModel):
    """Configuration for message/send requests."""

    accepted_output_modes: list[str] | None = Field(default=None, alias="acceptedOutputModes")
    history_length: int | None = Field(default=None, alias="historyLength")
    blocking: bool = False
    push_notification_config: PushNotificationConfig | None = Field(
        default=None, alias="pushNotificationConfig"
    )


# ═══════════════════════════════════════════════════════════════════════════════
# JSON-RPC WRAPPERS
# ═══════════════════════════════════════════════════════════════════════════════


class JsonRpcRequest(A2AModel):
    """JSON-RPC 2.0 request."""

    jsonrpc: Literal["2.0"] = "2.0"
    id: str | int
    method: str
    params: dict[str, Any] | None = None


class JsonRpcError(A2AModel):
    """JSON-RPC 2.0 error object."""

    code: int
    message: str
    data: dict[str, Any] | None = None


class JsonRpcResponse(A2AModel):
    """JSON-RPC 2.0 response."""

    jsonrpc: Literal["2.0"] = "2.0"
    id: str | int | None = None
    result: dict[str, Any] | None = None
    error: JsonRpcError | None = None


# ═══════════════════════════════════════════════════════════════════════════════
# A2A-SPECIFIC ERROR CODES
# ═══════════════════════════════════════════════════════════════════════════════


class A2AErrorCode:
    """A2A protocol error codes."""

    # JSON-RPC standard errors
    JSON_PARSE_ERROR = -32700
    INVALID_REQUEST = -32600
    METHOD_NOT_FOUND = -32601
    INVALID_PARAMS = -32602
    INTERNAL_ERROR = -32603

    # A2A-specific errors
    TASK_NOT_FOUND = -32001
    TASK_NOT_CANCELABLE = -32002
    PUSH_NOTIFICATION_NOT_SUPPORTED = -32003
    UNSUPPORTED_OPERATION = -32004
    CONTENT_TYPE_NOT_SUPPORTED = -32005
    INVALID_AGENT_RESPONSE = -32006
    AUTHENTICATED_EXTENDED_CARD_NOT_CONFIGURED = -32007
    EXTENSION_SUPPORT_REQUIRED = -32008
    VERSION_NOT_SUPPORTED = -32009


# ═══════════════════════════════════════════════════════════════════════════════
# SPECIFIC REQUEST/RESPONSE TYPES
# ═══════════════════════════════════════════════════════════════════════════════


class SendMessageParams(A2AModel):
    """Parameters for message/send."""

    message: Message
    configuration: MessageSendConfiguration | None = None
    metadata: dict[str, Any] | None = None


class GetTaskParams(A2AModel):
    """Parameters for tasks/get."""

    id: str
    history_length: int | None = Field(default=None, alias="historyLength")
    metadata: dict[str, Any] | None = None


class CancelTaskParams(A2AModel):
    """Parameters for tasks/cancel."""

    id: str
    metadata: dict[str, Any] | None = None


class ResubscribeParams(A2AModel):
    """Parameters for tasks/resubscribe."""

    id: str
    metadata: dict[str, Any] | None = None


class SendMessageResult(A2AModel):
    """Result from message/send - either a Task or Message."""

    task: Task | None = None
    message: Message | None = None


class StreamResponseResult(A2AModel):
    """Result from streaming responses."""

    task: Task | None = None
    message: Message | None = None
    status_update: TaskStatusUpdateEvent | None = Field(default=None, alias="statusUpdate")
    artifact_update: TaskArtifactUpdateEvent | None = Field(default=None, alias="artifactUpdate")


# ═══════════════════════════════════════════════════════════════════════════════
# PAGINATION
# ═══════════════════════════════════════════════════════════════════════════════


class ListTasksParams(A2AModel):
    """Parameters for tasks/list."""

    cursor: str | None = None
    page_size: int | None = Field(default=None, alias="pageSize", le=100)
    metadata: dict[str, Any] | None = None


class ListTasksResult(A2AModel):
    """Result from tasks/list."""

    tasks: list[Task]
    next_cursor: str | None = Field(default=None, alias="nextCursor")


# ═══════════════════════════════════════════════════════════════════════════════
# HELPER FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════════════


def is_terminal_state(state: TaskState | str) -> bool:
    """Check if a task state is terminal (no further transitions)."""
    if isinstance(state, str):
        try:
            state = TaskState(state)
        except ValueError:
            return False
    return state in {
        TaskState.COMPLETED,
        TaskState.FAILED,
        TaskState.CANCELED,
        TaskState.REJECTED,
    }


def is_interrupted_state(state: TaskState | str) -> bool:
    """Check if a task state is interrupted (waiting for input/auth)."""
    if isinstance(state, str):
        try:
            state = TaskState(state)
        except ValueError:
            return False
    return state in {TaskState.INPUT_REQUIRED, TaskState.AUTH_REQUIRED}


def create_text_message(
    text: str,
    role: MessageRole = MessageRole.USER,
    message_id: str | None = None,
    context_id: str | None = None,
    task_id: str | None = None,
) -> Message:
    """Create a simple text message."""
    import uuid

    return Message(
        role=role,
        messageId=message_id or str(uuid.uuid4()),
        parts=[TextPart(text=text)],
        contextId=context_id,
        taskId=task_id,
    )


def create_jsonrpc_request(
    method: str,
    params: dict[str, Any] | None = None,
    request_id: str | int | None = None,
) -> JsonRpcRequest:
    """Create a JSON-RPC 2.0 request."""
    import uuid

    return JsonRpcRequest(
        id=request_id or str(uuid.uuid4()),
        method=method,
        params=params,
    )
