a2a-lf 0.2.4

// Copyright AGNTCY Contributors (https://github.com/agntcy)
// SPDX-License-Identifier: Apache-2.0
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;

// ---------------------------------------------------------------------------
// TaskID / ArtifactID / ContextID
// ---------------------------------------------------------------------------

/// Unique identifier for a task, generated by the server.
pub type TaskId = String;

/// Unique identifier for an artifact within a task.
pub type ArtifactId = String;

/// Generate a new random task ID (UUIDv7).
pub fn new_task_id() -> TaskId {
    Uuid::now_v7().to_string()
}

/// Generate a new random context ID (UUIDv7).
pub fn new_context_id() -> String {
    Uuid::now_v7().to_string()
}

/// Generate a new random message ID (UUIDv7).
pub fn new_message_id() -> String {
    Uuid::now_v7().to_string()
}

/// Generate a new random artifact ID (UUIDv7).
pub fn new_artifact_id() -> ArtifactId {
    Uuid::now_v7().to_string()
}

// ---------------------------------------------------------------------------
// MessageRole
// ---------------------------------------------------------------------------

/// Identifies the sender of a message.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub enum Role {
    #[default]
    Unspecified,
    User,
    Agent,
}

impl Serialize for Role {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        match self {
            Role::Unspecified => serializer.serialize_str("ROLE_UNSPECIFIED"),
            Role::User => serializer.serialize_str("ROLE_USER"),
            Role::Agent => serializer.serialize_str("ROLE_AGENT"),
        }
    }
}

impl<'de> Deserialize<'de> for Role {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let s = String::deserialize(deserializer)?;
        match s.as_str() {
            "ROLE_USER" => Ok(Role::User),
            "ROLE_AGENT" => Ok(Role::Agent),
            "ROLE_UNSPECIFIED" | "" => Ok(Role::Unspecified),
            other => Err(serde::de::Error::unknown_variant(
                other,
                &["ROLE_USER", "ROLE_AGENT", "ROLE_UNSPECIFIED"],
            )),
        }
    }
}

// ---------------------------------------------------------------------------
// TaskState
// ---------------------------------------------------------------------------

/// The lifecycle state of a task.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub enum TaskState {
    #[default]
    Unspecified,
    Submitted,
    Working,
    Completed,
    Failed,
    Canceled,
    InputRequired,
    Rejected,
    AuthRequired,
}

impl TaskState {
    /// Returns true for terminal states where no further changes are permitted.
    pub fn is_terminal(&self) -> bool {
        matches!(
            self,
            TaskState::Completed | TaskState::Failed | TaskState::Canceled | TaskState::Rejected
        )
    }
}

impl Serialize for TaskState {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let s = match self {
            TaskState::Unspecified => "TASK_STATE_UNSPECIFIED",
            TaskState::Submitted => "TASK_STATE_SUBMITTED",
            TaskState::Working => "TASK_STATE_WORKING",
            TaskState::Completed => "TASK_STATE_COMPLETED",
            TaskState::Failed => "TASK_STATE_FAILED",
            TaskState::Canceled => "TASK_STATE_CANCELED",
            TaskState::InputRequired => "TASK_STATE_INPUT_REQUIRED",
            TaskState::Rejected => "TASK_STATE_REJECTED",
            TaskState::AuthRequired => "TASK_STATE_AUTH_REQUIRED",
        };
        serializer.serialize_str(s)
    }
}

impl<'de> Deserialize<'de> for TaskState {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let s = String::deserialize(deserializer)?;
        match s.as_str() {
            "TASK_STATE_SUBMITTED" => Ok(TaskState::Submitted),
            "TASK_STATE_WORKING" => Ok(TaskState::Working),
            "TASK_STATE_COMPLETED" => Ok(TaskState::Completed),
            "TASK_STATE_FAILED" => Ok(TaskState::Failed),
            "TASK_STATE_CANCELED" => Ok(TaskState::Canceled),
            "TASK_STATE_INPUT_REQUIRED" => Ok(TaskState::InputRequired),
            "TASK_STATE_REJECTED" => Ok(TaskState::Rejected),
            "TASK_STATE_AUTH_REQUIRED" => Ok(TaskState::AuthRequired),
            "TASK_STATE_UNSPECIFIED" | "" => Ok(TaskState::Unspecified),
            other => Err(serde::de::Error::unknown_variant(
                other,
                &[
                    "TASK_STATE_SUBMITTED",
                    "TASK_STATE_WORKING",
                    "TASK_STATE_COMPLETED",
                    "TASK_STATE_FAILED",
                    "TASK_STATE_CANCELED",
                    "TASK_STATE_INPUT_REQUIRED",
                    "TASK_STATE_REJECTED",
                    "TASK_STATE_AUTH_REQUIRED",
                    "TASK_STATE_UNSPECIFIED",
                ],
            )),
        }
    }
}

// ---------------------------------------------------------------------------
// Part (field-presence union)
// ---------------------------------------------------------------------------

/// A content part of a message or artifact. Uses field-presence serialization
/// (like C# SDK) so only the relevant content field is present in JSON.
#[derive(Debug, Clone, PartialEq)]
pub struct Part {
    pub content: PartContent,
    #[doc = "Optional filename for file parts."]
    pub filename: Option<String>,
    #[doc = "MIME type of the part content."]
    pub media_type: Option<String>,
    #[doc = "Optional metadata for extensions."]
    pub metadata: Option<HashMap<String, Value>>,
}

/// The content of a part — a discriminated union.
#[derive(Debug, Clone, PartialEq)]
pub enum PartContent {
    Text(String),
    Raw(Vec<u8>),
    Url(String),
    Data(Value),
}

impl Part {
    pub fn text(text: impl Into<String>) -> Self {
        Part {
            content: PartContent::Text(text.into()),
            filename: None,
            media_type: None,
            metadata: None,
        }
    }

    pub fn raw(data: Vec<u8>) -> Self {
        Part {
            content: PartContent::Raw(data),
            filename: None,
            media_type: None,
            metadata: None,
        }
    }

    pub fn url(url: impl Into<String>) -> Self {
        Part {
            content: PartContent::Url(url.into()),
            filename: None,
            media_type: None,
            metadata: None,
        }
    }

    pub fn data(value: Value) -> Self {
        Part {
            content: PartContent::Data(value),
            filename: None,
            media_type: None,
            metadata: None,
        }
    }

    pub fn with_media_type(mut self, media_type: impl Into<String>) -> Self {
        self.media_type = Some(media_type.into());
        self
    }

    pub fn with_filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = Some(filename.into());
        self
    }

    pub fn as_text(&self) -> Option<&str> {
        if let PartContent::Text(ref s) = self.content {
            Some(s)
        } else {
            None
        }
    }
}

impl Serialize for Part {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeMap;
        let mut map = serializer.serialize_map(None)?;

        match &self.content {
            PartContent::Text(t) => map.serialize_entry("text", t)?,
            PartContent::Raw(r) => map.serialize_entry("raw", &BASE64.encode(r))?,
            PartContent::Url(u) => map.serialize_entry("url", u)?,
            PartContent::Data(d) => map.serialize_entry("data", d)?,
        }

        if let Some(ref f) = self.filename {
            map.serialize_entry("filename", f)?;
        }
        if let Some(ref m) = self.media_type {
            map.serialize_entry("mediaType", m)?;
        }
        if let Some(ref meta) = self.metadata {
            if !meta.is_empty() {
                map.serialize_entry("metadata", meta)?;
            }
        }
        map.end()
    }
}

impl<'de> Deserialize<'de> for Part {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let raw: HashMap<String, Value> = HashMap::deserialize(deserializer)?;

        let content = if let Some(Value::String(t)) = raw.get("text") {
            PartContent::Text(t.clone())
        } else if let Some(Value::String(r)) = raw.get("raw") {
            let bytes = BASE64.decode(r).map_err(serde::de::Error::custom)?;
            PartContent::Raw(bytes)
        } else if let Some(Value::String(u)) = raw.get("url") {
            PartContent::Url(u.clone())
        } else if let Some(d) = raw.get("data") {
            PartContent::Data(d.clone())
        } else {
            return Err(serde::de::Error::custom(
                "Part must have one of: text, raw, url, data",
            ));
        };

        let filename = raw
            .get("filename")
            .and_then(|v| v.as_str())
            .map(String::from);
        let media_type = raw
            .get("mediaType")
            .and_then(|v| v.as_str())
            .map(String::from);
        let metadata: Option<HashMap<String, Value>> = raw
            .get("metadata")
            .and_then(|v| serde_json::from_value(v.clone()).ok());

        Ok(Part {
            content,
            filename,
            media_type,
            metadata,
        })
    }
}

// ---------------------------------------------------------------------------
// Message
// ---------------------------------------------------------------------------

/// A single message in a conversation between user and agent.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Message {
    /// Unique identifier for the message, generated by the sender.
    pub message_id: String,

    /// Context identifier for grouping related interactions.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_id: Option<String>,

    /// Task identifier this message belongs to.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub task_id: Option<TaskId>,

    /// Sender role.
    pub role: Role,

    /// Content parts forming the message body.
    pub parts: Vec<Part>,

    /// Optional metadata for extensions.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,

    /// Extension URIs relevant to this message.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub extensions: Option<Vec<String>>,

    /// Related task IDs for additional context.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reference_task_ids: Option<Vec<TaskId>>,
}

impl Message {
    /// Create a new message with a random ID.
    pub fn new(role: Role, parts: Vec<Part>) -> Self {
        Message {
            message_id: new_message_id(),
            context_id: None,
            task_id: None,
            role,
            parts,
            metadata: None,
            extensions: None,
            reference_task_ids: None,
        }
    }

    /// Return the text of the first text part, if any.
    pub fn text(&self) -> Option<&str> {
        self.parts.iter().find_map(|p| p.as_text())
    }
}

// ---------------------------------------------------------------------------
// TaskStatus
// ---------------------------------------------------------------------------

/// The status of a task at a specific point in time.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TaskStatus {
    pub state: TaskState,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub message: Option<Message>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub timestamp: Option<DateTime<Utc>>,
}

// ---------------------------------------------------------------------------
// Task
// ---------------------------------------------------------------------------

/// A single stateful operation or conversation between client and agent.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Task {
    pub id: TaskId,
    pub context_id: String,
    pub status: TaskStatus,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub artifacts: Option<Vec<Artifact>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub history: Option<Vec<Message>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,
}

// ---------------------------------------------------------------------------
// Artifact
// ---------------------------------------------------------------------------

/// A resource generated by an agent during a task.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Artifact {
    pub artifact_id: ArtifactId,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    pub parts: Vec<Part>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub extensions: Option<Vec<String>>,
}

// ---------------------------------------------------------------------------
// Request / Response types
// ---------------------------------------------------------------------------

/// Configuration for `SendMessage` / `SendStreamingMessage`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SendMessageConfiguration {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub accepted_output_modes: Option<Vec<String>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[serde(rename = "pushNotificationConfig")]
    pub push_notification_config: Option<PushNotificationConfig>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub history_length: Option<i32>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub return_immediately: Option<bool>,
}

/// Request to send a message to an agent.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SendMessageRequest {
    pub message: Message,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub configuration: Option<SendMessageConfiguration>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Response from `SendMessage` — either a Task or a Message.
#[derive(Debug, Clone, PartialEq)]
pub enum SendMessageResponse {
    Task(Task),
    Message(Message),
}

impl Serialize for SendMessageResponse {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeMap;
        let mut map = serializer.serialize_map(Some(1))?;
        match self {
            SendMessageResponse::Task(t) => map.serialize_entry("task", t)?,
            SendMessageResponse::Message(m) => map.serialize_entry("message", m)?,
        }
        map.end()
    }
}

impl<'de> Deserialize<'de> for SendMessageResponse {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let raw: HashMap<String, Value> = HashMap::deserialize(deserializer)?;
        if let Some(v) = raw.get("task") {
            let task: Task = serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?;
            Ok(SendMessageResponse::Task(task))
        } else if let Some(v) = raw.get("message") {
            let msg: Message =
                serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?;
            Ok(SendMessageResponse::Message(msg))
        } else {
            Err(serde::de::Error::custom(
                "SendMessageResponse must have 'task' or 'message'",
            ))
        }
    }
}

/// Request to get a task by ID.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetTaskRequest {
    pub id: TaskId,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub history_length: Option<i32>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to list tasks.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListTasksRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_id: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub status: Option<TaskState>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub page_size: Option<i32>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub page_token: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub history_length: Option<i32>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub status_timestamp_after: Option<DateTime<Utc>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub include_artifacts: Option<bool>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Response for listing tasks.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListTasksResponse {
    pub tasks: Vec<Task>,
    pub next_page_token: String,
    pub page_size: i32,
    pub total_size: i32,
}

/// Request to cancel a task.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelTaskRequest {
    pub id: TaskId,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, Value>>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to subscribe to task updates.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SubscribeToTaskRequest {
    pub id: TaskId,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to get an extended (authenticated) agent card.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetExtendedAgentCardRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

// ---------------------------------------------------------------------------
// Push Notification types
// ---------------------------------------------------------------------------

/// Push notification configuration.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PushNotificationConfig {
    /// Callback URL for push notifications.
    pub url: String,

    /// Optional unique ID for this configuration.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,

    /// Optional session/task token for validation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub token: Option<String>,

    /// Optional authentication details for the webhook.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub authentication: Option<AuthenticationInfo>,
}

/// Authentication details for a push notification endpoint.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AuthenticationInfo {
    /// HTTP authentication scheme (e.g., "Bearer", "Basic").
    pub scheme: String,

    /// Scheme-specific credentials.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub credentials: Option<String>,
}

/// Container associating push config with a task.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TaskPushNotificationConfig {
    pub task_id: TaskId,
    pub config: PushNotificationConfig,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to get a specific push config.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetTaskPushNotificationConfigRequest {
    pub task_id: TaskId,
    pub id: String,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to list push configs for a task.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListTaskPushNotificationConfigsRequest {
    pub task_id: TaskId,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub page_size: Option<i32>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub page_token: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Response for listing push configs.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListTaskPushNotificationConfigsResponse {
    pub configs: Vec<TaskPushNotificationConfig>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next_page_token: Option<String>,
}

/// Request to create a push config for a task.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateTaskPushNotificationConfigRequest {
    pub task_id: TaskId,
    pub config: PushNotificationConfig,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

/// Request to delete a push config.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DeleteTaskPushNotificationConfigRequest {
    pub task_id: TaskId,
    pub id: String,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tenant: Option<String>,
}

// ---------------------------------------------------------------------------
// TransportProtocol
// ---------------------------------------------------------------------------

/// Transport protocol identifier. Open-form string to allow custom protocols.
pub type TransportProtocol = String;

/// JSON-RPC transport protocol constant.
pub const TRANSPORT_PROTOCOL_JSONRPC: &str = "JSONRPC";
/// gRPC transport protocol constant.
pub const TRANSPORT_PROTOCOL_GRPC: &str = "GRPC";
/// HTTP+JSON (REST) transport protocol constant.
pub const TRANSPORT_PROTOCOL_HTTP_JSON: &str = "HTTP+JSON";
/// SLIMRPC transport protocol constant.
pub const TRANSPORT_PROTOCOL_SLIMRPC: &str = "SLIMRPC";

/// Protocol version string (e.g., "1.0").
pub type ProtocolVersion = String;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_task_state_serde() {
        let state = TaskState::Completed;
        let json = serde_json::to_string(&state).unwrap();
        assert_eq!(json, r#""TASK_STATE_COMPLETED""#);
        let back: TaskState = serde_json::from_str(&json).unwrap();
        assert_eq!(back, TaskState::Completed);
    }

    #[test]
    fn test_role_serde() {
        let role = Role::Agent;
        let json = serde_json::to_string(&role).unwrap();
        assert_eq!(json, r#""ROLE_AGENT""#);
        let back: Role = serde_json::from_str(&json).unwrap();
        assert_eq!(back, Role::Agent);
    }

    #[test]
    fn test_part_text_serde() {
        let part = Part::text("hello");
        let json = serde_json::to_string(&part).unwrap();
        let v: Value = serde_json::from_str(&json).unwrap();
        assert_eq!(v["text"], "hello");
        assert!(v.get("raw").is_none());

        let back: Part = serde_json::from_str(&json).unwrap();
        assert_eq!(back.content, PartContent::Text("hello".into()));
    }

    #[test]
    fn test_part_raw_serde() {
        let part = Part::raw(vec![1, 2, 3]);
        let json = serde_json::to_string(&part).unwrap();
        let back: Part = serde_json::from_str(&json).unwrap();
        assert_eq!(back.content, PartContent::Raw(vec![1, 2, 3]));
    }

    #[test]
    fn test_send_message_response_serde() {
        let resp = SendMessageResponse::Task(Task {
            id: "t1".into(),
            context_id: "c1".into(),
            status: TaskStatus {
                state: TaskState::Submitted,
                message: None,
                timestamp: None,
            },
            artifacts: None,
            history: None,
            metadata: None,
        });
        let json = serde_json::to_string(&resp).unwrap();
        let v: Value = serde_json::from_str(&json).unwrap();
        assert!(v.get("task").is_some());
        assert!(v.get("message").is_none());
        let back: SendMessageResponse = serde_json::from_str(&json).unwrap();
        assert!(matches!(back, SendMessageResponse::Task(_)));
    }

    #[test]
    fn test_send_message_response_message_variant() {
        let resp = SendMessageResponse::Message(Message::new(Role::Agent, vec![Part::text("hi")]));
        let json = serde_json::to_string(&resp).unwrap();
        let v: Value = serde_json::from_str(&json).unwrap();
        assert!(v.get("message").is_some());
        assert!(v.get("task").is_none());
        let back: SendMessageResponse = serde_json::from_str(&json).unwrap();
        assert!(matches!(back, SendMessageResponse::Message(_)));
    }

    #[test]
    fn test_part_url_serde() {
        let part = Part::url("https://example.com/file.pdf");
        let json = serde_json::to_string(&part).unwrap();
        let v: Value = serde_json::from_str(&json).unwrap();
        assert_eq!(v["url"], "https://example.com/file.pdf");
        let back: Part = serde_json::from_str(&json).unwrap();
        assert!(matches!(back.content, PartContent::Url(_)));
    }

    #[test]
    fn test_part_data_serde() {
        let data = serde_json::json!({"key": "value", "count": 42});
        let part = Part::data(data.clone());
        let json = serde_json::to_string(&part).unwrap();
        let back: Part = serde_json::from_str(&json).unwrap();
        assert!(matches!(back.content, PartContent::Data(_)));
        if let PartContent::Data(d) = back.content {
            assert_eq!(d["key"], "value");
        }
    }

    #[test]
    fn test_part_with_metadata_and_media_type() {
        let part = Part::text("hello")
            .with_media_type("text/plain")
            .with_filename("test.txt");
        assert_eq!(part.media_type.as_deref(), Some("text/plain"));
        assert_eq!(part.filename.as_deref(), Some("test.txt"));
        let json = serde_json::to_string(&part).unwrap();
        let v: Value = serde_json::from_str(&json).unwrap();
        assert_eq!(v["mediaType"], "text/plain");
        assert_eq!(v["filename"], "test.txt");
    }

    #[test]
    fn test_part_as_text() {
        let part = Part::text("hello");
        assert_eq!(part.as_text(), Some("hello"));
        let part = Part::raw(vec![]);
        assert_eq!(part.as_text(), None);
    }

    #[test]
    fn test_message_new() {
        let msg = Message::new(Role::User, vec![Part::text("hi")]);
        assert!(!msg.message_id.is_empty());
        assert_eq!(msg.role, Role::User);
        assert_eq!(msg.parts.len(), 1);
        assert_eq!(msg.text(), Some("hi"));
    }

    #[test]
    fn test_message_serde() {
        let msg = Message {
            message_id: "m1".to_string(),
            context_id: Some("c1".to_string()),
            task_id: None,
            role: Role::Agent,
            parts: vec![Part::text("response")],
            metadata: None,
            extensions: Some(vec!["ext1".to_string()]),
            reference_task_ids: Some(vec!["t1".to_string()]),
        };
        let json = serde_json::to_string(&msg).unwrap();
        let back: Message = serde_json::from_str(&json).unwrap();
        assert_eq!(msg, back);
    }

    #[test]
    fn test_task_state_is_terminal() {
        assert!(TaskState::Completed.is_terminal());
        assert!(TaskState::Failed.is_terminal());
        assert!(TaskState::Canceled.is_terminal());
        assert!(TaskState::Rejected.is_terminal());
        assert!(!TaskState::Submitted.is_terminal());
        assert!(!TaskState::Working.is_terminal());
        assert!(!TaskState::InputRequired.is_terminal());
        assert!(!TaskState::AuthRequired.is_terminal());
        assert!(!TaskState::Unspecified.is_terminal());
    }

    #[test]
    fn test_task_full_serde() {
        let task = Task {
            id: "t1".to_string(),
            context_id: "c1".to_string(),
            status: TaskStatus {
                state: TaskState::Working,
                message: Some(Message::new(Role::Agent, vec![Part::text("working")])),
                timestamp: None,
            },
            artifacts: Some(vec![Artifact {
                artifact_id: "a1".to_string(),
                name: Some("output".to_string()),
                description: None,
                parts: vec![Part::text("result data")],
                metadata: None,
                extensions: None,
            }]),
            history: Some(vec![Message::new(Role::User, vec![Part::text("do it")])]),
            metadata: None,
        };
        let json = serde_json::to_string(&task).unwrap();
        let back: Task = serde_json::from_str(&json).unwrap();
        assert_eq!(task.id, back.id);
        assert_eq!(task.status.state, back.status.state);
        assert!(back.artifacts.is_some());
        assert!(back.history.is_some());
    }

    #[test]
    fn test_push_notification_config_serde() {
        let config = PushNotificationConfig {
            url: "https://example.com/webhook".to_string(),
            id: Some("cfg-1".to_string()),
            token: Some("tok-1".to_string()),
            authentication: Some(AuthenticationInfo {
                scheme: "Bearer".to_string(),
                credentials: Some("secret".to_string()),
            }),
        };
        let json = serde_json::to_string(&config).unwrap();
        let back: PushNotificationConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(config, back);
    }

    #[test]
    fn test_send_message_request_serde() {
        let req = SendMessageRequest {
            message: Message::new(Role::User, vec![Part::text("hello")]),
            configuration: Some(SendMessageConfiguration {
                accepted_output_modes: Some(vec!["text/plain".to_string()]),
                push_notification_config: None,
                history_length: Some(10),
                return_immediately: Some(true),
            }),
            metadata: None,
            tenant: Some("tenant-1".to_string()),
        };
        let json = serde_json::to_string(&req).unwrap();
        let back: SendMessageRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(req.tenant, back.tenant);
        assert_eq!(
            req.configuration.as_ref().unwrap().history_length,
            back.configuration.as_ref().unwrap().history_length
        );
    }

    #[test]
    fn test_role_all_variants_serde() {
        let cases = [
            (Role::Unspecified, "\"ROLE_UNSPECIFIED\""),
            (Role::User, "\"ROLE_USER\""),
            (Role::Agent, "\"ROLE_AGENT\""),
        ];

        for (role, expected_json) in cases {
            let json = serde_json::to_string(&role).unwrap();
            assert_eq!(json, expected_json);
            let back: Role = serde_json::from_str(&json).unwrap();
            assert_eq!(back, role);
        }

        let back: Role = serde_json::from_str("\"\"").unwrap();
        assert_eq!(back, Role::Unspecified);
    }

    #[test]
    fn test_task_state_all_variants_serde() {
        let cases = [
            (TaskState::Unspecified, "TASK_STATE_UNSPECIFIED"),
            (TaskState::Submitted, "TASK_STATE_SUBMITTED"),
            (TaskState::Working, "TASK_STATE_WORKING"),
            (TaskState::Completed, "TASK_STATE_COMPLETED"),
            (TaskState::Failed, "TASK_STATE_FAILED"),
            (TaskState::Canceled, "TASK_STATE_CANCELED"),
            (TaskState::InputRequired, "TASK_STATE_INPUT_REQUIRED"),
            (TaskState::Rejected, "TASK_STATE_REJECTED"),
            (TaskState::AuthRequired, "TASK_STATE_AUTH_REQUIRED"),
        ];

        for (state, expected_str) in cases {
            let json = serde_json::to_string(&state).unwrap();
            assert_eq!(json, format!("\"{}\"", expected_str));
            let back: TaskState = serde_json::from_str(&json).unwrap();
            assert_eq!(back, state);
        }

        let back: TaskState = serde_json::from_str("\"\"").unwrap();
        assert_eq!(back, TaskState::Unspecified);
    }

    #[test]
    fn test_part_all_content_variants_serde() {
        let parts = [
            Part::text("hello"),
            Part::raw(vec![1, 2, 3]),
            Part::url("https://example.com"),
            Part::data(serde_json::json!({"value": 1})),
        ];

        for part in parts {
            let json = serde_json::to_string(&part).unwrap();
            let back: Part = serde_json::from_str(&json).unwrap();
            match (&part.content, &back.content) {
                (PartContent::Text(a), PartContent::Text(b)) => assert_eq!(a, b),
                (PartContent::Raw(a), PartContent::Raw(b)) => assert_eq!(a, b),
                (PartContent::Url(a), PartContent::Url(b)) => assert_eq!(a, b),
                (PartContent::Data(_), PartContent::Data(_)) => {}
                _ => panic!("mismatched part content variants"),
            }
        }
    }

    #[test]
    fn test_send_message_response_message_deserialize() {
        let json = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_AGENT",
                "parts": [{"text": "hello"}]
            }
        });
        let back: SendMessageResponse = serde_json::from_value(json).unwrap();
        assert!(matches!(back, SendMessageResponse::Message(_)));
    }

    #[test]
    fn test_new_id_functions() {
        let task_id = new_task_id();
        assert!(!task_id.is_empty());
        let ctx_id = new_context_id();
        assert!(!ctx_id.is_empty());
        let msg_id = new_message_id();
        assert!(!msg_id.is_empty());
        let art_id = new_artifact_id();
        assert!(!art_id.is_empty());
        // All should be unique
        assert_ne!(task_id, ctx_id);
    }

    #[test]
    fn test_role_default() {
        assert_eq!(Role::default(), Role::Unspecified);
    }

    #[test]
    fn test_task_state_default() {
        assert_eq!(TaskState::default(), TaskState::Unspecified);
    }

    #[test]
    fn test_list_tasks_request_serde() {
        let req = ListTasksRequest {
            context_id: Some("c1".to_string()),
            status: Some(TaskState::Working),
            page_size: Some(10),
            page_token: None,
            history_length: Some(5),
            status_timestamp_after: None,
            include_artifacts: Some(true),
            tenant: None,
        };
        let json = serde_json::to_string(&req).unwrap();
        let back: ListTasksRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(req, back);
    }

    #[test]
    fn test_cancel_task_request_serde() {
        let req = CancelTaskRequest {
            id: "t1".to_string(),
            metadata: None,
            tenant: Some("ten".to_string()),
        };
        let json = serde_json::to_string(&req).unwrap();
        let back: CancelTaskRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(req, back);
    }

    #[test]
    fn test_subscribe_to_task_request_serde() {
        let req = SubscribeToTaskRequest {
            id: "t1".to_string(),
            tenant: None,
        };
        let json = serde_json::to_string(&req).unwrap();
        let back: SubscribeToTaskRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(req, back);
    }
}