use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TaskState {
    Waiting,
    Working,
    Completed,
    Failed,
}

impl TaskState {
    pub fn as_str(&self) -> &'static str {
        match self {
            TaskState::Waiting => "waiting",
            TaskState::Working => "working",
            TaskState::Completed => "completed",
            TaskState::Failed => "failed",
        }
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct A2aTask {
    pub id: String,
    pub message: A2aMessage,
    #[serde(default)]
    pub metadata: HashMap<String, Value>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct A2aMessage {
    pub role: String,
    pub parts: Vec<A2aMessagePart>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type", content = "value")]
pub enum A2aMessagePart {
    #[serde(rename = "text")]
    Text(String),
    #[serde(rename = "file")]
    File {
        path: String,
        mime_type: Option<String>,
    },
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct A2aTaskStatus {
    pub id: String,
    pub state: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<A2aMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<A2aTaskResult>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<A2aErrorObj>,
    pub created_at: String,
    pub updated_at: String,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct A2aTaskResult {
    pub message: A2aMessage,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub artifacts: Option<Vec<A2aArtifact>>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct A2aArtifact {
    #[serde(rename = "type")]
    pub artifact_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub format: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub title: Option<String>,
    pub data: Value,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct A2aErrorObj {
    pub code: i32,
    pub message: String,
}

#[derive(Debug, Deserialize)]
pub struct JsonRpcRequest<T> {
    pub jsonrpc: String,
    pub id: Value,
    pub method: String,
    pub params: T,
}

#[derive(Debug, Serialize)]
pub struct JsonRpcResponse<T> {
    pub jsonrpc: String,
    pub id: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<T>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcErrorObj>,
}

#[derive(Debug, Serialize)]
pub struct JsonRpcErrorObj {
    pub code: i32,
    pub message: String,
}

impl<T: Serialize> JsonRpcResponse<T> {
    pub fn success(id: Value, result: T) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: Some(result),
            error: None,
        }
    }

    pub fn error(id: Value, code: i32, message: String) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: None,
            error: Some(JsonRpcErrorObj { code, message }),
        }
    }
}
