pmcp 2.2.0

//! Types for composition operations.

use serde::{Deserialize, Serialize};

/// Content from a resource read operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceContent {
    /// The URI of the resource.
    pub uri: String,

    /// The MIME type of the content, if known.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mime_type: Option<String>,

    /// Text content, if the resource contains text.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,

    /// Binary content as base64, if the resource contains binary data.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub blob: Option<String>,
}

impl ResourceContent {
    /// Create a new text resource content.
    pub fn text(uri: impl Into<String>, text: impl Into<String>) -> Self {
        Self {
            uri: uri.into(),
            mime_type: Some("text/plain".to_string()),
            text: Some(text.into()),
            blob: None,
        }
    }

    /// Create a new JSON resource content.
    pub fn json(uri: impl Into<String>, text: impl Into<String>) -> Self {
        Self {
            uri: uri.into(),
            mime_type: Some("application/json".to_string()),
            text: Some(text.into()),
            blob: None,
        }
    }

    /// Create a new binary resource content.
    pub fn binary(
        uri: impl Into<String>,
        mime_type: impl Into<String>,
        blob: impl Into<String>,
    ) -> Self {
        Self {
            uri: uri.into(),
            mime_type: Some(mime_type.into()),
            text: None,
            blob: Some(blob.into()),
        }
    }

    /// Check if the content is text.
    pub fn is_text(&self) -> bool {
        self.text.is_some()
    }

    /// Check if the content is binary.
    pub fn is_binary(&self) -> bool {
        self.blob.is_some()
    }

    /// Get text content, if available.
    pub fn as_text(&self) -> Option<&str> {
        self.text.as_deref()
    }
}

/// Result from a prompt operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptResult {
    /// Description of the prompt.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Messages generated by the prompt.
    pub messages: Vec<PromptMessage>,
}

impl PromptResult {
    /// Create a new prompt result with messages.
    pub fn new(messages: Vec<PromptMessage>) -> Self {
        Self {
            description: None,
            messages,
        }
    }

    /// Create a new prompt result with a description.
    pub fn with_description(description: impl Into<String>, messages: Vec<PromptMessage>) -> Self {
        Self {
            description: Some(description.into()),
            messages,
        }
    }

    /// Get the first message, if any.
    pub fn first_message(&self) -> Option<&PromptMessage> {
        self.messages.first()
    }

    /// Get all user messages.
    pub fn user_messages(&self) -> impl Iterator<Item = &PromptMessage> {
        self.messages.iter().filter(|m| m.role == "user")
    }

    /// Get all assistant messages.
    pub fn assistant_messages(&self) -> impl Iterator<Item = &PromptMessage> {
        self.messages.iter().filter(|m| m.role == "assistant")
    }
}

/// A message in a prompt result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptMessage {
    /// The role of the message (e.g., "user", "assistant", "system").
    pub role: String,

    /// The content of the message.
    pub content: PromptContent,
}

impl PromptMessage {
    /// Create a new user message with text content.
    pub fn user(text: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: PromptContent::Text { text: text.into() },
        }
    }

    /// Create a new assistant message with text content.
    pub fn assistant(text: impl Into<String>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: PromptContent::Text { text: text.into() },
        }
    }

    /// Create a new system message with text content.
    pub fn system(text: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: PromptContent::Text { text: text.into() },
        }
    }

    /// Get the text content, if this is a text message.
    pub fn as_text(&self) -> Option<&str> {
        match &self.content {
            PromptContent::Text { text } => Some(text),
            _ => None,
        }
    }
}

/// Content of a prompt message.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum PromptContent {
    /// Text content.
    Text {
        /// The text content.
        text: String,
    },

    /// Image content.
    Image {
        /// Base64-encoded image data.
        data: String,
        /// MIME type of the image.
        #[serde(rename = "mimeType")]
        mime_type: String,
    },

    /// Resource content (embedded resource).
    Resource {
        /// Resource reference.
        resource: EmbeddedResource,
    },
}

/// An embedded resource in a prompt message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddedResource {
    /// URI of the resource.
    pub uri: String,

    /// MIME type of the resource.
    #[serde(skip_serializing_if = "Option::is_none", rename = "mimeType")]
    pub mime_type: Option<String>,

    /// Text content if available.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,

    /// Binary content as base64 if available.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub blob: Option<String>,
}

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

    #[test]
    fn test_resource_content_text() {
        let content = ResourceContent::text("file://test.txt", "Hello, world!");
        assert!(content.is_text());
        assert!(!content.is_binary());
        assert_eq!(content.as_text(), Some("Hello, world!"));
    }

    #[test]
    fn test_resource_content_json() {
        let content = ResourceContent::json("file://data.json", r#"{"key": "value"}"#);
        assert!(content.is_text());
        assert_eq!(content.mime_type, Some("application/json".to_string()));
    }

    #[test]
    fn test_prompt_message_user() {
        let msg = PromptMessage::user("Hello!");
        assert_eq!(msg.role, "user");
        assert_eq!(msg.as_text(), Some("Hello!"));
    }

    #[test]
    fn test_prompt_result() {
        let result = PromptResult::with_description(
            "Test prompt",
            vec![
                PromptMessage::user("What is 2+2?"),
                PromptMessage::assistant("4"),
            ],
        );

        assert_eq!(result.description, Some("Test prompt".to_string()));
        assert_eq!(result.messages.len(), 2);
        assert_eq!(result.user_messages().count(), 1);
        assert_eq!(result.assistant_messages().count(), 1);
    }

    #[test]
    fn test_prompt_content_serialization() {
        let content = PromptContent::Text {
            text: "Hello".to_string(),
        };
        let json = serde_json::to_string(&content).unwrap();
        assert!(json.contains(r#""type":"text""#));
        assert!(json.contains(r#""text":"Hello""#));
    }
}