# ruff: noqa: E501
from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, Field, field_validator


class Task(BaseModel):
    """Represents the stateful unit of work being processed by the A2A Server for an
    A2A Client. A task encapsulates the entire interaction related to a specific goal or
    request.
    """

    id: str = Field(description="Unique identifier for the task")
    context_id: str = Field(
        alias="contextId",
        description="Server-generated id for contextual alignment across interactions",
    )
    status: TaskStatus = Field(description="Current status of the task")
    history: list[Message] | None = None
    artifacts: list[Artifact] | None = Field(
        None, description="Collection of artifacts created by the agent."
    )
    metadata: dict[str, Any] | None = Field(None, description="Extension metadata.")
    kind: Literal["task"] = Field("task", description="Event type")


class TaskStatus(BaseModel):
    """TaskState and accompanying message.
    Represents the current state and associated context (e.g., a message from the
    agent) of a Task."""

    state: TaskState
    message: Message | None = Field(
        None, description="Additional status updates for client"
    )
    timestamp: str | None = Field(
        None,
        description="ISO 8601 datetime string when the status was recorded.",
        examples=["2023-10-27T10:00:00Z"],
    )


class TaskState(str, Enum):
    """Represents the possible states of a Task.

    Values:
        SUBMITTED: Task received by the server and acknowledged, but processing has not yet actively started.
        WORKING: Task is actively being processed by the agent. Client may expect further updates or a terminal state.
        INPUT_REQUIRED: Agent requires additional input from the client/user to proceed. The task is effectively paused.
        COMPLETED: Task finished successfully. Results are typically available in Task.artifacts or TaskStatus.message.
        CANCELED: Task was canceled (e.g., by a tasks/cancel request or server-side policy).
        FAILED: Task terminated due to an error during processing. TaskStatus.message may contain error details.
        REJECTED: Task terminated due to rejection by remote agent. TaskStatus.message may contain error details.
        AUTH_REQUIRED: Agent requires additional authentication from the client/user to proceed. The task is effectively paused.
        UNKNOWN: The state of the task cannot be determined (e.g., task ID is invalid, unknown, or has expired).
    """

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


class Artifact(BaseModel):
    """Represents a tangible output generated by the agent during a task. Artifacts
    are the results or products of the agent's work."""

    artifact_id: str = Field(
        alias="artifactId",
        description="Identifier for the artifact generated by the agent",
    )
    name: str | None = Field(None, description="Descriptive name for the artifact")
    description: str | None = Field(
        None, description="Human-readable description of the artifact"
    )
    parts: list[Part] = Field(
        description="Content of the artifact, as one or more Part objects. Must have at least one"
    )
    metadata: dict[str, Any] | None = Field(
        None, description="Arbitrary key-value metadata associated with the artifact"
    )
    extensions: list[str] | None = Field(
        None, description="A list of extension URIs that contributed to this artifact"
    )


class Message(BaseModel):
    """Represents a single communication turn or a piece of contextual information
    between a client and an agent. Messages are used for instructions, prompts,
    replies, and status updates."""

    role: Literal["user", "agent"]
    parts: list[Part] = Field(
        description="Message content. Array of content parts. Must contain at least one part."
    )
    metadata: dict[str, Any] | None = None
    reference_task_ids: list[str] | None = Field(None, alias="referenceTaskIds")
    message_id: str = Field(alias="messageId")
    task_id: str | None = Field(None, alias="taskId")
    context_id: str | None = Field(None, alias="contextId")
    kind: Literal["message"] = Field("message", description="Event type")


class PartBase(BaseModel):
    """Base properties common to all message parts."""

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


class TextPart(PartBase):
    """Represents a text segment within parts."""

    kind: Literal["text"]
    text: str


class FilePart(PartBase):
    """Represents a File segment within parts."""

    kind: Literal["file"]
    file: FileWithBytes | FileWithUri


class FileBase(BaseModel):
    """Represents the base entity for FileParts."""

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


class FileWithBytes(FileBase):
    """Define the variant where 'bytes' is present and 'uri' is absent."""

    bytes: str
    uri: None = None


class FileWithUri(FileBase):
    """Define the variant where 'uri' is present and 'bytes' is absent."""

    uri: str
    bytes: None = None


class DataPart(PartBase):
    """Represents a structured data segment within a message part."""

    kind: Literal["data"]
    data: dict[str, Any]


Part = TextPart | FilePart | DataPart


class PushNotificationAuthenticationInfo(BaseModel):
    """Defines authentication details for push notifications."""

    schemes: list[str] = Field(
        description="Supported authentication schemes - e.g. Basic, Bearer"
    )
    credentials: str | None = Field(None, description="Optional credentials")


class PushNotificationConfig(BaseModel):
    """Configuration for setting up push notifications for task updates."""

    id: str | None = Field(
        None,
        description="Push Notification ID - created by server to support multiple callbacks",
    )
    url: str = Field(description="URL for sending the push notifications")
    token: str | None = Field(None, description="Token unique to this task/session")
    authentication: PushNotificationAuthenticationInfo | None = None


class TaskPushNotificationConfig(BaseModel):
    """Parameters for setting or getting push notification configuration for a task."""

    task_id: str = Field(alias="taskId", description="Task id")
    push_notification_config: PushNotificationConfig = Field(
        alias="pushNotificationConfig", description="Push notification configuration"
    )


# ----------
# SecurityScheme, ..., PasswordOAuthFlow
# ----------


# ----------
# Request and Response Models
# ----------


class JSONRPCMessage(BaseModel):
    """Base JSON-RPC 2.0 message."""

    jsonrpc: Literal["2.0"] = Field(
        description="Specifies the version of the JSON-RPC protocol. MUST be exactly '2.0'."
    )
    id: int | str | None = Field(
        None,
        description="An identifier established by the Client that MUST contain a String, Number. Numbers SHOULD NOT contain fractional parts.",
    )


class JSONRPCRequest(JSONRPCMessage):
    """Represents a JSON-RPC 2.0 Request object."""

    method: str = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: dict[str, Any] | None = Field(
        None,
        description="A Structured value that holds the parameter values to be used during the invocation of the method.",
    )


class JSONRPCError(BaseModel):
    """Represents a JSON-RPC 2.0 Error object.
    This is typically included in a JSONRPCErrorResponse when an error occurs.
    """

    code: int = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(
        description="A String providing a short description of the error."
    )
    data: Any | None = Field(
        None,
        description="A Primitive or Structured value that contains additional information about the error. This may be omitted.",
    )

    @field_validator("code", mode="after")
    @classmethod
    def validate_code(cls, v: int) -> int:
        if not (-32000 <= v <= -32099):
            raise ValueError(f"Invalid code: {v}", "code must be in [-32000, -32099]")
        return v


class JSONRPCSuccessResponse(JSONRPCMessage):
    """Represents a JSON-RPC 2.0 Success Response object."""

    id: int | str | None = Field(description="Request identifier")
    result: Any = Field(description="The result object on success")


class JSONRPCErrorResponse(JSONRPCMessage):
    """Represents a JSON-RPC 2.0 Error Response object."""

    id: int | str | None = Field(description="Request identifier")
    error: JSONRPCError | A2AError = Field(description="Error details")


# ----------
# "message/send"
# Sends a message to an agent to initiate a new interaction or to continue an existing one.
# This method is suitable for synchronous request/response interactions or when client-side
# polling (using tasks/get) is acceptable for monitoring longer-running tasks.
# ----------


class MessageSendConfiguration(BaseModel):
    """Configuration for the send message request."""

    accepted_output_modes: list[str] = Field(
        alias="acceptedOutputModes",
        description="Accepted output modalities by the client.",
    )
    history_length: int | None = Field(
        None,
        alias="historyLength",
        description="Number of recent messages to be retrieved.",
    )
    push_notification_config: PushNotificationConfig | None = Field(
        None,
        alias="pushNotificationConfig",
        description="Where the server should send notifications when disconnected.",
    )
    blocking: bool | None = Field(
        None, description="If the server should treat the client as a blocking request."
    )


class MessageSendParams(BaseModel):
    """Sent by the client to the agent as a request. May create, continue or restart a task."""

    message: Message = Field(description="The message being sent to the server.")
    configuration: MessageSendConfiguration | None = Field(
        None, description="Send message configuration."
    )
    metadata: dict[str, Any] | None = Field(None, description="Extension metadata.")


class SendMessageRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'message/send' method."""

    id: int | str
    method: Literal["message/send"]
    params: MessageSendParams


class SendMessageSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response model for the 'message/send' method."""

    result: Message | Task


SendMessageResponse = SendMessageSuccessResponse | JSONRPCErrorResponse


# ----------
# "message/stream"
# ----------


class TaskStatusUpdateEvent(BaseModel):
    """Sent by server during sendStream or subscribe requests."""

    task_id: str = Field(alias="taskId", description="Task id")
    context_id: str = Field(
        alias="contextId", description="The context the task is associated with"
    )
    kind: Literal["status-update"] = Field(description="Event type")
    status: TaskStatus = Field(description="Current status of the task")
    final: bool = Field(description="Indicates the end of the event stream")
    metadata: dict[str, Any] | None = Field(None, description="Extension metadata.")


class TaskArtifactUpdateEvent(BaseModel):
    """Sent by server during sendStream or subscribe requests."""

    task_id: str = Field(alias="taskId", description="Task id")
    context_id: str = Field(
        alias="contextId", description="The context the task is associated with"
    )
    kind: Literal["artifact-update"] = Field(description="Event type")
    artifact: Artifact = Field(description="Generated artifact")
    append: bool | None = Field(
        None, description="Indicates if this artifact appends to a previous one"
    )
    last_chunk: bool | None = Field(
        None,
        alias="lastChunk",
        description="Indicates if this is the last chunk of the artifact",
    )
    metadata: dict[str, Any] | None = Field(None, description="Extension metadata.")


class SendStreamingMessageRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'message/stream' method."""

    id: int | str
    method: Literal["message/stream"]
    params: MessageSendParams


class SendStreamingMessageSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response model for the 'message/stream' method."""

    result: Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent


SendStreamingMessageResponse = (
    SendStreamingMessageSuccessResponse | JSONRPCErrorResponse
)


# ----------
# "tasks/get"
# Retrieves the current state (including status, artifacts, and optionally history) of a
# previously initiated task. This is typically used for polling the status of a task
# initiated with message/send, or for fetching the final state of a task after being
# notified via a push notification or after an SSE stream has ended.
# ----------


class TaskIdParams(BaseModel):
    """Parameters containing only a task ID, used for simple task operations."""

    id: str = Field(description="Task id.")
    metadata: dict[str, Any] | None = None


class TaskQueryParams(TaskIdParams):
    """Parameters for querying a task, including optional history length."""

    history_length: int | None = Field(
        None,
        alias="historyLength",
        description="Number of recent messages to be retrieved.",
    )


class GetTaskRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'tasks/get' method."""

    id: int | str
    method: Literal["tasks/get"] = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: TaskQueryParams = Field(
        description="A Structured value that holds the parameter values to be used during the invocation of the method."
    )


class GetTaskSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response for the 'tasks/get' method."""

    result: Task = Field(description="The result object on success.")


GetTaskResponse = GetTaskSuccessResponse | JSONRPCErrorResponse


# ----------
# "tasks/cancel"
# ----------


class CancelTaskRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'tasks/cancel' method."""

    id: int | str
    method: Literal["tasks/cancel"] = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: TaskIdParams = Field(
        description="A Structured value that holds the parameter values to be used during the invocation of the method."
    )


class CancelTaskSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response model for the 'tasks/cancel' method."""

    result: Task = Field(description="The result object on success.")


CancelTaskResponse = CancelTaskSuccessResponse | JSONRPCErrorResponse


# ----------
# "tasks/pushNotificationConfig/set"
# ----------


class SetTaskPushNotificationConfigRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'tasks/pushNotificationConfig/set' method."""

    id: int | str
    method: Literal["tasks/pushNotificationConfig/set"] = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: TaskPushNotificationConfig = Field(
        description="A Structured value that holds the parameter values to be used during the invocation of the method."
    )


class SetTaskPushNotificationConfigSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response model for the 'tasks/pushNotificationConfig/set' method."""

    result: TaskPushNotificationConfig = Field(
        description="The result object on success."
    )


SetTaskPushNotificationConfigResponse = (
    SetTaskPushNotificationConfigSuccessResponse | JSONRPCErrorResponse
)


# ----------
# "tasks/pushNotificationConfig/get"
# ----------


class GetTaskPushNotificationConfigRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'tasks/pushNotificationConfig/get' method."""

    id: int | str
    method: Literal["tasks/pushNotificationConfig/get"] = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: TaskIdParams = Field(
        description="A Structured value that holds the parameter values to be used during the invocation of the method."
    )


class GetTaskPushNotificationConfigSuccessResponse(JSONRPCSuccessResponse):
    """JSON-RPC success response model for the 'tasks/pushNotificationConfig/get' method."""

    result: TaskPushNotificationConfig = Field(
        description="The result object on success."
    )


GetTaskPushNotificationConfigResponse = (
    GetTaskPushNotificationConfigSuccessResponse | JSONRPCErrorResponse
)


# ----------
# "tasks/resubscribe"
# ----------


class TaskResubscriptionRequest(JSONRPCRequest):
    """JSON-RPC request model for the 'tasks/resubscribe' method."""

    id: int | str
    method: Literal["tasks/resubscribe"] = Field(
        description="A String containing the name of the method to be invoked."
    )
    params: TaskIdParams = Field(
        description="A Structured value that holds the parameter values to be used during the invocation of the method."
    )


# ----------
# JSON-RPC Response
# ----------


JSONRPCResponse = (
    SendMessageResponse
    | SendStreamingMessageResponse
    | GetTaskResponse
    | CancelTaskResponse
    | SetTaskPushNotificationConfigResponse
    | GetTaskPushNotificationConfigResponse
)

# ----------
# A2A
# ----------


A2ARequest = (
    SendMessageRequest
    | SendStreamingMessageRequest
    | GetTaskRequest
    | CancelTaskRequest
    | SetTaskPushNotificationConfigRequest
    | GetTaskPushNotificationConfigRequest
    | TaskResubscriptionRequest
)


# ----------
# Error Models
# ----------


class JSONParseError(JSONRPCError):
    """JSON-RPC error indicating invalid JSON was received by the server."""

    code: Literal[-32700]
    message: str = Field(default="Invalid JSON payload")


class InvalidRequestError(JSONRPCError):
    """JSON-RPC error indicating the JSON sent is not a valid Request object."""

    code: Literal[-32600] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Request payload validation error")


class MethodNotFoundError(JSONRPCError):
    """JSON-RPC error indicating the method does not exist or is not available."""

    code: Literal[-32601] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Method not found")


class InvalidParamsError(JSONRPCError):
    """JSON-RPC error indicating invalid method parameter(s)."""

    code: Literal[-32602] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Invalid parameters")


class InternalError(JSONRPCError):
    """JSON-RPC error indicating an internal JSON-RPC error on the server."""

    code: Literal[-32603] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Internal error")


class TaskNotFoundError(JSONRPCError):
    """A2A specific error indicating the requested task ID was not found."""

    code: Literal[-32001] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Task not found")


class TaskNotCancelableError(JSONRPCError):
    """A2A specific error indicating the task is in a state where it cannot be canceled."""

    code: Literal[-32002] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Task cannot be canceled")


class PushNotificationNotSupportedError(JSONRPCError):
    """A2A specific error indicating the agent does not support push notifications."""

    code: Literal[-32003] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Push Notification is not supported")


class UnsupportedOperationError(JSONRPCError):
    """A2A specific error indicating the requested operation is not supported by the agent."""

    code: Literal[-32004] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="This operation is not supported")


class ContentTypeNotSupportedError(JSONRPCError):
    """A2A specific error indicating incompatible content types between request and agent capabilities."""

    code: Literal[-32005] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Incompatible content types")


class InvalidAgentResponseError(JSONRPCError):
    """A2A specific error indicating agent returned invalid response for the current method."""

    code: Literal[-32006] = Field(
        description="A Number that indicates the error type that occurred."
    )
    message: str = Field(default="Invalid agent response")


A2AError = (
    JSONParseError
    | InvalidRequestError
    | MethodNotFoundError
    | InvalidParamsError
    | InternalError
    | TaskNotFoundError
    | TaskNotCancelableError
    | PushNotificationNotSupportedError
    | UnsupportedOperationError
    | ContentTypeNotSupportedError
    | InvalidAgentResponseError
)
