coderlib 0.1.0

//! Storage module for CoderLib
//!
//! This module provides abstractions and implementations for persisting
//! sessions, messages, and other data.

pub mod sqlite;
pub mod memory;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::core::error::StorageError;
use crate::core::session::Session;

pub use sqlite::SqliteStorage;
pub use memory::MemoryStorage;

/// Trait for storage backends
#[async_trait]
pub trait Storage: Send + Sync {
    /// Create a new session
    async fn create_session(&self, session: &Session) -> Result<String, StorageError>;
    
    /// Get a session by ID
    async fn get_session(&self, id: &str) -> Result<Option<Session>, StorageError>;
    
    /// Update an existing session
    async fn update_session(&self, session: &Session) -> Result<(), StorageError>;
    
    /// Delete a session
    async fn delete_session(&self, id: &str) -> Result<(), StorageError>;
    
    /// List all sessions
    async fn list_sessions(&self) -> Result<Vec<Session>, StorageError>;
    
    /// Save a message
    async fn save_message(&self, message: &Message) -> Result<String, StorageError>;
    
    /// Get messages for a session
    async fn get_messages(&self, session_id: &str) -> Result<Vec<Message>, StorageError>;
    
    /// Delete a message
    async fn delete_message(&self, id: &str) -> Result<(), StorageError>;
    
    /// Search messages by content
    async fn search_messages(&self, query: &str) -> Result<Vec<Message>, StorageError>;
    
    /// Get storage statistics
    async fn get_statistics(&self) -> Result<StorageStatistics, StorageError>;
    
    /// Perform cleanup operations (e.g., delete old sessions)
    async fn cleanup(&self, older_than: DateTime<Utc>) -> Result<u64, StorageError>;
}

/// A message in a conversation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// Unique message identifier
    pub id: String,
    
    /// Session this message belongs to
    pub session_id: String,
    
    /// Role of the message sender
    pub role: MessageRole,
    
    /// Message content
    pub content: MessageContent,
    
    /// When the message was created
    pub timestamp: DateTime<Utc>,
    
    /// Additional metadata
    pub metadata: serde_json::Value,
}

/// Role of a message sender
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum MessageRole {
    /// Message from the user
    User,
    
    /// Message from the AI assistant
    Assistant,
    
    /// System message (prompts, instructions)
    System,
    
    /// Tool execution result
    Tool,
}

/// Content of a message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MessageContent {
    /// Plain text content
    Text(String),

    /// Structured content with multiple parts
    Structured(Vec<ContentPart>),

    /// Tool call content
    ToolCall {
        name: String,
        parameters: serde_json::Value,
        id: String,
    },

    /// Tool result content
    ToolResult {
        result: String,
        tool_call_id: String,
    },
}

/// Part of structured message content
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ContentPart {
    /// Text content
    Text {
        text: String,
    },
    
    /// Image content
    Image {
        /// Image data (base64 encoded)
        data: String,
        /// MIME type
        mime_type: String,
        /// Optional description
        description: Option<String>,
    },
    
    /// File attachment
    File {
        /// File path
        path: String,
        /// File content (base64 encoded)
        content: String,
        /// MIME type
        mime_type: String,
    },
    
    /// Code block
    Code {
        /// Programming language
        language: String,
        /// Code content
        code: String,
    },
    
    /// Tool call
    ToolCall {
        /// Tool name
        name: String,
        /// Tool parameters
        parameters: serde_json::Value,
    },
    
    /// Tool result
    ToolResult {
        /// Tool call ID this result corresponds to
        call_id: String,
        /// Result content
        result: String,
        /// Whether the tool execution was successful
        success: bool,
    },
}

/// Storage statistics
#[derive(Debug, Clone)]
pub struct StorageStatistics {
    /// Total number of sessions
    pub total_sessions: u64,
    
    /// Total number of messages
    pub total_messages: u64,
    
    /// Storage size in bytes
    pub storage_size: u64,
    
    /// Oldest session timestamp
    pub oldest_session: Option<DateTime<Utc>>,
    
    /// Newest session timestamp
    pub newest_session: Option<DateTime<Utc>>,
    
    /// Average messages per session
    pub avg_messages_per_session: f64,
}

impl MessageContent {
    /// Get the text content of the message
    pub fn as_text(&self) -> String {
        match self {
            MessageContent::Text(text) => text.clone(),
            MessageContent::Structured(parts) => {
                parts
                    .iter()
                    .filter_map(|part| match part {
                        ContentPart::Text { text } => Some(text.clone()),
                        ContentPart::Code { code, .. } => Some(code.clone()),
                        _ => None,
                    })
                    .collect::<Vec<_>>()
                    .join("\n")
            }
            MessageContent::ToolCall { name, parameters, .. } => {
                format!("Tool call: {} with parameters: {}", name, parameters)
            }
            MessageContent::ToolResult { result, .. } => result.clone(),
        }
    }
    
    /// Check if the message contains any images
    pub fn has_images(&self) -> bool {
        match self {
            MessageContent::Text(_) => false,
            MessageContent::Structured(parts) => {
                parts.iter().any(|part| matches!(part, ContentPart::Image { .. }))
            }
            MessageContent::ToolCall { .. } | MessageContent::ToolResult { .. } => false,
        }
    }
    
    /// Check if the message contains any files
    pub fn has_files(&self) -> bool {
        match self {
            MessageContent::Text(_) => false,
            MessageContent::Structured(parts) => {
                parts.iter().any(|part| matches!(part, ContentPart::File { .. }))
            }
            MessageContent::ToolCall { .. } | MessageContent::ToolResult { .. } => false,
        }
    }
    
    /// Check if the message contains any tool calls
    pub fn has_tool_calls(&self) -> bool {
        match self {
            MessageContent::Text(_) => false,
            MessageContent::Structured(parts) => {
                parts.iter().any(|part| matches!(part, ContentPart::ToolCall { .. }))
            }
            MessageContent::ToolCall { .. } => true,
            MessageContent::ToolResult { .. } => false,
        }
    }
    
    /// Get all tool calls in the message
    pub fn get_tool_calls(&self) -> Vec<&ContentPart> {
        match self {
            MessageContent::Text(_) => Vec::new(),
            MessageContent::Structured(parts) => {
                parts
                    .iter()
                    .filter(|part| matches!(part, ContentPart::ToolCall { .. }))
                    .collect()
            }
            MessageContent::ToolCall { .. } | MessageContent::ToolResult { .. } => Vec::new(),
        }
    }
}

impl std::fmt::Display for MessageRole {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MessageRole::User => write!(f, "user"),
            MessageRole::Assistant => write!(f, "assistant"),
            MessageRole::System => write!(f, "system"),
            MessageRole::Tool => write!(f, "tool"),
        }
    }
}

impl std::str::FromStr for MessageRole {
    type Err = String;
    
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "user" => Ok(MessageRole::User),
            "assistant" => Ok(MessageRole::Assistant),
            "system" => Ok(MessageRole::System),
            "tool" => Ok(MessageRole::Tool),
            _ => Err(format!("Unknown message role: {}", s)),
        }
    }
}

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

    #[test]
    fn test_message_content_as_text() {
        let text_content = MessageContent::Text("Hello, world!".to_string());
        assert_eq!(text_content.as_text(), "Hello, world!");
        
        let structured_content = MessageContent::Structured(vec![
            ContentPart::Text { text: "Hello".to_string() },
            ContentPart::Code { 
                language: "rust".to_string(), 
                code: "println!(\"world\");".to_string() 
            },
        ]);
        assert_eq!(structured_content.as_text(), "Hello\nprintln!(\"world\");");
    }

    #[test]
    fn test_message_content_has_images() {
        let text_content = MessageContent::Text("Hello".to_string());
        assert!(!text_content.has_images());
        
        let image_content = MessageContent::Structured(vec![
            ContentPart::Text { text: "Look at this:".to_string() },
            ContentPart::Image { 
                data: "base64data".to_string(),
                mime_type: "image/png".to_string(),
                description: None,
            },
        ]);
        assert!(image_content.has_images());
    }

    #[test]
    fn test_message_role_display() {
        assert_eq!(MessageRole::User.to_string(), "user");
        assert_eq!(MessageRole::Assistant.to_string(), "assistant");
        assert_eq!(MessageRole::System.to_string(), "system");
        assert_eq!(MessageRole::Tool.to_string(), "tool");
    }

    #[test]
    fn test_message_role_from_str() {
        assert_eq!("user".parse::<MessageRole>().unwrap(), MessageRole::User);
        assert_eq!("ASSISTANT".parse::<MessageRole>().unwrap(), MessageRole::Assistant);
        assert!("invalid".parse::<MessageRole>().is_err());
    }

    #[test]
    fn test_message_content_tool_calls() {
        let content = MessageContent::Structured(vec![
            ContentPart::Text { text: "I'll help you with that.".to_string() },
            ContentPart::ToolCall {
                name: "file_read".to_string(),
                parameters: serde_json::json!({"path": "test.txt"}),
            },
        ]);
        
        assert!(content.has_tool_calls());
        assert_eq!(content.get_tool_calls().len(), 1);
    }
}