//! A2A Protocol request and response types
//!
//! Implements the wire format for A2A invoke/stream/status/cancel endpoints.

use crate::message::A2aMessageType;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A2A task status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum A2aTaskStatus {
    Pending,
    Running,
    Completed,
    Failed,
    Cancelled,
}

impl A2aTaskStatus {
    pub fn is_terminal(&self) -> bool {
        matches!(self, Self::Completed | Self::Failed | Self::Cancelled)
    }
}

/// A2A request payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct A2aMessagePayload {
    pub message_type: A2aMessageType,
    pub action: String,
    pub data: serde_json::Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub correlation_id: Option<String>,
}

/// A2A invoke/stream request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct A2aRequest {
    pub task_id: String,
    pub from_agent: String,
    pub to_agent: String,
    pub message: A2aMessagePayload,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

impl A2aRequest {
    /// Create a new A2A request
    pub fn new(
        task_id: impl Into<String>,
        from_agent: impl Into<String>,
        to_agent: impl Into<String>,
        message_type: A2aMessageType,
        action: impl Into<String>,
        data: serde_json::Value,
    ) -> Self {
        Self {
            task_id: task_id.into(),
            from_agent: from_agent.into(),
            to_agent: to_agent.into(),
            message: A2aMessagePayload {
                message_type,
                action: action.into(),
                data,
                correlation_id: None,
            },
            metadata: None,
        }
    }

    /// Builder: set correlation ID
    pub fn with_correlation_id(mut self, id: impl Into<String>) -> Self {
        self.message.correlation_id = Some(id.into());
        self
    }

    /// Builder: set metadata
    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
        self.metadata = Some(metadata);
        self
    }
}

/// A2A response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct A2aResponse {
    pub task_id: String,
    pub status: A2aTaskStatus,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

impl A2aResponse {
    /// Create a pending response
    pub fn pending(task_id: impl Into<String>) -> Self {
        Self {
            task_id: task_id.into(),
            status: A2aTaskStatus::Pending,
            result: None,
            error: None,
        }
    }

    /// Create a running response
    pub fn running(task_id: impl Into<String>) -> Self {
        Self {
            task_id: task_id.into(),
            status: A2aTaskStatus::Running,
            result: None,
            error: None,
        }
    }

    /// Create a completed response
    pub fn completed(task_id: impl Into<String>, result: serde_json::Value) -> Self {
        Self {
            task_id: task_id.into(),
            status: A2aTaskStatus::Completed,
            result: Some(result),
            error: None,
        }
    }

    /// Create a failed response
    pub fn failed(task_id: impl Into<String>, error: impl Into<String>) -> Self {
        Self {
            task_id: task_id.into(),
            status: A2aTaskStatus::Failed,
            result: None,
            error: Some(error.into()),
        }
    }

    /// Create a cancelled response
    pub fn cancelled(task_id: impl Into<String>) -> Self {
        Self {
            task_id: task_id.into(),
            status: A2aTaskStatus::Cancelled,
            result: None,
            error: None,
        }
    }
}

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

    #[test]
    fn test_request_creation() {
        let req = A2aRequest::new(
            "task-123",
            "agent-a",
            "agent-b",
            A2aMessageType::ProcessData,
            "analyze",
            json!({"text": "hello"}),
        );
        assert_eq!(req.task_id, "task-123");
        assert_eq!(req.message.message_type, A2aMessageType::ProcessData);
    }

    #[test]
    fn test_request_serialization() {
        let req = A2aRequest::new(
            "task-123",
            "agent-a",
            "agent-b",
            A2aMessageType::ExecuteAction,
            "run",
            json!({}),
        );
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("EXECUTE_ACTION"));
        assert!(json.contains("task-123"));
    }

    #[test]
    fn test_response_helpers() {
        let pending = A2aResponse::pending("t1");
        assert_eq!(pending.status, A2aTaskStatus::Pending);
        assert!(!pending.status.is_terminal());

        let completed = A2aResponse::completed("t2", json!({"output": "done"}));
        assert_eq!(completed.status, A2aTaskStatus::Completed);
        assert!(completed.status.is_terminal());

        let failed = A2aResponse::failed("t3", "error occurred");
        assert_eq!(failed.status, A2aTaskStatus::Failed);
        assert!(failed.error.is_some());
    }

    #[test]
    fn test_task_status_serialization() {
        let status = A2aTaskStatus::Running;
        let json = serde_json::to_string(&status).unwrap();
        assert_eq!(json, "\"running\"");
    }
}
