coderlib 0.1.0

//! Conversation summarizer for generating conversation summaries

use tracing::{debug, info, warn};

use crate::llm::Provider;
use crate::storage::{Message, MessageRole, MessageContent};
use super::{SummarizationConfig, SummarizationError, ConversationSummary, utils};

/// Conversation summarizer that generates summaries using an LLM
pub struct ConversationSummarizer {
    /// Provider for summarization
    provider: Box<dyn Provider>,
    
    /// Configuration
    config: SummarizationConfig,
}

impl ConversationSummarizer {
    /// Create a new conversation summarizer
    pub fn new(provider: Box<dyn Provider>, config: SummarizationConfig) -> Self {
        Self { provider, config }
    }
    
    /// Summarize a conversation from messages
    pub async fn summarize(
        &self,
        session_id: &str,
        messages: Vec<Message>,
    ) -> Result<ConversationSummary, SummarizationError> {
        if !self.config.enabled {
            return Err(SummarizationError::Disabled);
        }
        
        if messages.is_empty() {
            return Err(SummarizationError::EmptyConversation);
        }
        
        info!("Starting summarization for session {} with {} messages", session_id, messages.len());
        
        // Prepare messages for summarization
        let (messages_to_summarize, preserved_messages) = self.prepare_messages(messages);
        
        if messages_to_summarize.is_empty() {
            return Err(SummarizationError::EmptyConversation);
        }
        
        // Add summarization prompt
        let mut summarization_messages = messages_to_summarize;
        summarization_messages.push(Message {
            id: uuid::Uuid::new_v4().to_string(),
            session_id: session_id.to_string(),
            role: MessageRole::User,
            content: MessageContent::Text(utils::create_summary_prompt()),
            timestamp: chrono::Utc::now(),
            metadata: serde_json::Value::Null,
        });
        
        debug!("Sending {} messages to summarizer", summarization_messages.len());
        
        // Generate summary using the provider
        let response = self.provider
            .send_messages(summarization_messages, vec![])
            .await
            .map_err(|e| SummarizationError::ProviderError(e.to_string()))?;
        
        let summary_text = response.content.trim().to_string();
        
        if summary_text.is_empty() {
            return Err(SummarizationError::SummaryGenerationFailed(
                "Empty summary returned from provider".to_string()
            ));
        }
        
        let summary_tokens = response.token_usage
            .map(|u| u.output_tokens)
            .unwrap_or_else(|| utils::estimate_token_count(&summary_text));
        
        info!("Generated summary for session {} ({} tokens)", session_id, summary_tokens);
        
        Ok(ConversationSummary {
            original_session_id: session_id.to_string(),
            summary_text,
            preserved_messages,
            summary_tokens,
            summarizer_model: self.provider.model().id.clone(),
            created_at: chrono::Utc::now(),
        })
    }
    
    /// Prepare messages for summarization by filtering and preserving recent ones
    fn prepare_messages(&self, mut messages: Vec<Message>) -> (Vec<Message>, Vec<Message>) {
        // Sort messages by timestamp to ensure proper order
        messages.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
        
        let preserve_count = self.config.preserve_recent_messages as usize;
        
        if messages.len() <= preserve_count {
            // If we have fewer messages than we want to preserve, summarize all but preserve none
            return (messages, vec![]);
        }
        
        // Split messages: most for summarization, recent ones for preservation
        let split_point = messages.len() - preserve_count;
        let messages_to_summarize = messages[..split_point].to_vec();
        let preserved_messages = messages[split_point..].to_vec();
        
        debug!(
            "Prepared {} messages for summarization, preserving {} recent messages",
            messages_to_summarize.len(),
            preserved_messages.len()
        );
        
        (messages_to_summarize, preserved_messages)
    }
    
    /// Get the model being used for summarization
    pub fn model(&self) -> &crate::llm::Model {
        self.provider.model()
    }
    
    /// Get the configuration
    pub fn config(&self) -> &SummarizationConfig {
        &self.config
    }
    
    /// Check if the summarizer is available
    pub async fn is_available(&self) -> bool {
        self.provider.is_available().await
    }
    
    /// Estimate the token count for a set of messages
    pub fn estimate_message_tokens(&self, messages: &[Message]) -> u32 {
        messages.iter()
            .map(|msg| match &msg.content {
                MessageContent::Text(text) => utils::estimate_token_count(text),
                MessageContent::Structured(parts) => {
                    parts.iter().map(|part| match part {
                        crate::storage::ContentPart::Text { text } => utils::estimate_token_count(text),
                        crate::storage::ContentPart::Image { .. } => 1000, // Rough estimate for image processing
                        crate::storage::ContentPart::File { content, .. } => utils::estimate_token_count(content),
                        crate::storage::ContentPart::Code { code, .. } => utils::estimate_token_count(code),
                        crate::storage::ContentPart::ToolCall { parameters, .. } => {
                            utils::estimate_token_count(&parameters.to_string())
                        }
                        crate::storage::ContentPart::ToolResult { result, .. } => {
                            utils::estimate_token_count(result)
                        }
                    }).sum()
                }
                MessageContent::ToolCall { name, parameters, .. } => {
                    utils::estimate_token_count(name) + utils::estimate_token_count(&parameters.to_string())
                }
                MessageContent::ToolResult { result, .. } => utils::estimate_token_count(result),
            })
            .sum()
    }
    
    /// Validate that the messages can be summarized
    pub fn validate_messages(&self, messages: &[Message]) -> Result<(), SummarizationError> {
        if messages.is_empty() {
            return Err(SummarizationError::EmptyConversation);
        }
        
        // Check if we have at least one user or assistant message
        let has_content = messages.iter().any(|msg| {
            matches!(msg.role, MessageRole::User | MessageRole::Assistant)
        });
        
        if !has_content {
            return Err(SummarizationError::EmptyConversation);
        }
        
        // Estimate total tokens
        let estimated_tokens = self.estimate_message_tokens(messages);
        let model_context = utils::get_model_context_window(&self.provider.model().id);
        
        if estimated_tokens > model_context {
            warn!(
                "Message tokens ({}) exceed model context window ({}), summarization may fail",
                estimated_tokens, model_context
            );
        }
        
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::providers::MockProvider;
    use crate::llm::{Model, ModelCapabilities};
    
    fn create_test_messages(count: usize) -> Vec<Message> {
        (0..count)
            .map(|i| Message {
                id: format!("msg-{}", i),
                session_id: "test-session".to_string(),
                role: if i % 2 == 0 { MessageRole::User } else { MessageRole::Assistant },
                content: MessageContent::Text(format!("Test message {}", i)),
                timestamp: chrono::Utc::now(),
                metadata: serde_json::Value::Null,
            })
            .collect()
    }
    
    fn create_test_model() -> Model {
        Model {
            id: "test-model".to_string(),
            name: "Test Model".to_string(),
            provider: "test".to_string(),
            context_length: 4096,
            max_output_tokens: 2048,
            supports_tools: false,
            supports_streaming: false,
            supports_vision: false,
            cost_per_input_token: 0.0,
            cost_per_output_token: 0.0,
            capabilities: ModelCapabilities::default(),
        }
    }
    
    #[tokio::test]
    async fn test_prepare_messages() {
        let config = SummarizationConfig {
            preserve_recent_messages: 3,
            ..Default::default()
        };
        
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        let messages = create_test_messages(10);
        let (to_summarize, preserved) = summarizer.prepare_messages(messages);
        
        assert_eq!(to_summarize.len(), 7);
        assert_eq!(preserved.len(), 3);
    }
    
    #[tokio::test]
    async fn test_prepare_messages_few_messages() {
        let config = SummarizationConfig {
            preserve_recent_messages: 5,
            ..Default::default()
        };
        
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        let messages = create_test_messages(3);
        let (to_summarize, preserved) = summarizer.prepare_messages(messages);
        
        assert_eq!(to_summarize.len(), 3);
        assert_eq!(preserved.len(), 0);
    }
    
    #[tokio::test]
    async fn test_estimate_message_tokens() {
        let config = SummarizationConfig::default();
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        let messages = vec![
            Message {
                id: "1".to_string(),
                session_id: "test".to_string(),
                role: MessageRole::User,
                content: MessageContent::Text("Hello world".to_string()),
                timestamp: chrono::Utc::now(),
                metadata: serde_json::Value::Null,
            }
        ];
        
        let tokens = summarizer.estimate_message_tokens(&messages);
        assert!(tokens > 0);
        assert!(tokens < 100); // Should be reasonable for short text
    }
    
    #[tokio::test]
    async fn test_validate_messages() {
        let config = SummarizationConfig::default();
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        // Test empty messages
        let empty_messages = vec![];
        assert!(summarizer.validate_messages(&empty_messages).is_err());
        
        // Test valid messages
        let valid_messages = create_test_messages(5);
        assert!(summarizer.validate_messages(&valid_messages).is_ok());
    }
    
    #[tokio::test]
    async fn test_summarize_disabled() {
        let config = SummarizationConfig {
            enabled: false,
            ..Default::default()
        };
        
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        let messages = create_test_messages(5);
        let result = summarizer.summarize("test-session", messages).await;
        
        assert!(matches!(result, Err(SummarizationError::Disabled)));
    }
    
    #[tokio::test]
    async fn test_summarize_empty_conversation() {
        let config = SummarizationConfig::default();
        let model = create_test_model();
        let provider = Box::new(MockProvider::new());
        let summarizer = ConversationSummarizer::new(provider, config);
        
        let empty_messages = vec![];
        let result = summarizer.summarize("test-session", empty_messages).await;
        
        assert!(matches!(result, Err(SummarizationError::EmptyConversation)));
    }
}