coderlib 0.1.0

//! Token usage monitoring for auto-summarization

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

use crate::core::TokenUsage;
use super::{SummarizationConfig, SummarizationError, ConversationMetrics, utils};

/// Token usage monitor for tracking conversation metrics
pub struct TokenMonitor {
    /// Per-session token usage tracking
    session_usage: Arc<RwLock<HashMap<String, TokenUsage>>>,
    
    /// Model context window mappings
    model_context_windows: HashMap<String, u32>,
    
    /// Configuration
    config: SummarizationConfig,
}

impl TokenMonitor {
    /// Create a new token monitor
    pub fn new(config: SummarizationConfig) -> Self {
        let mut context_windows = HashMap::new();
        
        // Populate with known model context windows
        Self::populate_context_windows(&mut context_windows);
        
        Self {
            session_usage: Arc::new(RwLock::new(HashMap::new())),
            model_context_windows: context_windows,
            config,
        }
    }
    
    /// Populate the context windows map with known models
    fn populate_context_windows(context_windows: &mut HashMap<String, u32>) {
        // OpenAI models
        context_windows.insert("gpt-4".to_string(), 8192);
        context_windows.insert("gpt-4-0613".to_string(), 8192);
        context_windows.insert("gpt-4-32k".to_string(), 32768);
        context_windows.insert("gpt-4-32k-0613".to_string(), 32768);
        context_windows.insert("gpt-4-turbo".to_string(), 128000);
        context_windows.insert("gpt-4-turbo-preview".to_string(), 128000);
        context_windows.insert("gpt-4o".to_string(), 128000);
        context_windows.insert("gpt-4o-mini".to_string(), 128000);
        context_windows.insert("gpt-3.5-turbo".to_string(), 16385);
        context_windows.insert("gpt-3.5-turbo-0125".to_string(), 16385);
        
        // Anthropic models
        context_windows.insert("claude-3-5-sonnet-20241022".to_string(), 200000);
        context_windows.insert("claude-3-5-sonnet-20240620".to_string(), 200000);
        context_windows.insert("claude-3-opus-20240229".to_string(), 200000);
        context_windows.insert("claude-3-sonnet-20240229".to_string(), 200000);
        context_windows.insert("claude-3-haiku-20240307".to_string(), 200000);
        
        // Google models
        context_windows.insert("gemini-1.5-pro".to_string(), 2000000);
        context_windows.insert("gemini-1.5-pro-latest".to_string(), 2000000);
        context_windows.insert("gemini-1.5-flash".to_string(), 1000000);
        context_windows.insert("gemini-1.5-flash-latest".to_string(), 1000000);
        context_windows.insert("gemini-pro".to_string(), 32768);
        
        // Other common models
        context_windows.insert("llama-2-70b".to_string(), 4096);
        context_windows.insert("llama-2-13b".to_string(), 4096);
        context_windows.insert("llama-2-7b".to_string(), 4096);
        context_windows.insert("codellama-34b".to_string(), 16384);
        context_windows.insert("codellama-13b".to_string(), 16384);
        context_windows.insert("codellama-7b".to_string(), 16384);
        context_windows.insert("mistral-7b".to_string(), 32768);
        context_windows.insert("mixtral-8x7b".to_string(), 32768);
    }
    
    /// Update token usage for a session
    pub async fn update_usage(
        &self,
        session_id: &str,
        usage: TokenUsage,
    ) -> Result<(), SummarizationError> {
        let mut session_usage = self.session_usage.write().await;
        
        // Add to existing usage or create new entry
        let current_usage = session_usage.entry(session_id.to_string()).or_default();
        current_usage.input_tokens += usage.input_tokens;
        current_usage.output_tokens += usage.output_tokens;
        current_usage.cache_creation_tokens += usage.cache_creation_tokens;
        current_usage.cache_read_tokens += usage.cache_read_tokens;
        current_usage.update_total();
        
        debug!(
            "Updated token usage for session {}: {} total tokens",
            session_id, current_usage.total_tokens
        );
        
        Ok(())
    }
    
    /// Get current token usage for a session
    pub async fn get_usage(&self, session_id: &str) -> TokenUsage {
        let session_usage = self.session_usage.read().await;
        session_usage.get(session_id).cloned().unwrap_or_default()
    }
    
    /// Get conversation metrics for a session
    pub async fn get_metrics(
        &self,
        session_id: &str,
        model: &str,
        message_count: u32,
    ) -> Result<ConversationMetrics, SummarizationError> {
        let session_usage = self.session_usage.read().await;
        let usage = session_usage.get(session_id).cloned().unwrap_or_default();
        
        let context_window = self.get_context_window(model);
        let usage_percentage = utils::calculate_usage_percentage(usage.total_tokens, context_window);
        
        let metrics = ConversationMetrics {
            total_tokens: usage.total_tokens,
            message_count,
            context_window,
            usage_percentage,
            needs_summarization: utils::should_summarize(&ConversationMetrics {
                total_tokens: usage.total_tokens,
                message_count,
                context_window,
                usage_percentage,
                needs_summarization: false, // Will be set below
            }, &self.config),
        };
        
        debug!(
            "Session {} metrics: {} tokens ({:.1}% of {}), {} messages, needs_summarization: {}",
            session_id,
            metrics.total_tokens,
            metrics.usage_percentage * 100.0,
            metrics.context_window,
            metrics.message_count,
            metrics.needs_summarization
        );
        
        Ok(metrics)
    }
    
    /// Get context window size for a model
    pub fn get_context_window(&self, model: &str) -> u32 {
        self.model_context_windows
            .get(model)
            .copied()
            .unwrap_or_else(|| {
                // Try to infer from model name patterns
                let context_window = utils::get_model_context_window(model);
                warn!(
                    "Unknown model '{}', using inferred context window: {}",
                    model, context_window
                );
                context_window
            })
    }
    
    /// Add or update a model's context window
    pub fn register_model(&mut self, model: &str, context_window: u32) {
        self.model_context_windows.insert(model.to_string(), context_window);
        info!("Registered model '{}' with context window: {}", model, context_window);
    }
    
    /// Clear usage data for a session
    pub async fn clear_session(&self, session_id: &str) -> Result<(), SummarizationError> {
        let mut session_usage = self.session_usage.write().await;
        session_usage.remove(session_id);
        debug!("Cleared token usage for session: {}", session_id);
        Ok(())
    }
    
    /// Get all tracked sessions
    pub async fn get_tracked_sessions(&self) -> Vec<String> {
        let session_usage = self.session_usage.read().await;
        session_usage.keys().cloned().collect()
    }
    
    /// Get total token usage across all sessions
    pub async fn get_total_usage(&self) -> TokenUsage {
        let session_usage = self.session_usage.read().await;
        let mut total = TokenUsage::default();

        for usage in session_usage.values() {
            total.input_tokens += usage.input_tokens;
            total.output_tokens += usage.output_tokens;
            total.cache_creation_tokens += usage.cache_creation_tokens;
            total.cache_read_tokens += usage.cache_read_tokens;
        }
        total.update_total();
        
        total
    }
    
    /// Check if a session needs summarization
    pub async fn needs_summarization(
        &self,
        session_id: &str,
        model: &str,
        message_count: u32,
    ) -> Result<bool, SummarizationError> {
        let metrics = self.get_metrics(session_id, model, message_count).await?;
        Ok(metrics.needs_summarization)
    }
    
    /// Get usage statistics
    pub async fn get_usage_stats(&self) -> UsageStats {
        let session_usage = self.session_usage.read().await;
        let session_count = session_usage.len();
        let total_usage = session_usage.values().fold(TokenUsage::default(), |mut acc, usage| {
            acc.input_tokens += usage.input_tokens;
            acc.output_tokens += usage.output_tokens;
            acc.cache_creation_tokens += usage.cache_creation_tokens;
            acc.cache_read_tokens += usage.cache_read_tokens;
            acc.update_total();
            acc
        });
        
        let average_tokens_per_session = if session_count > 0 {
            total_usage.total_tokens / session_count as u32
        } else {
            0
        };
        
        UsageStats {
            session_count,
            total_usage,
            average_tokens_per_session,
        }
    }
}

/// Usage statistics
#[derive(Debug, Clone)]
pub struct UsageStats {
    /// Number of tracked sessions
    pub session_count: usize,
    
    /// Total token usage across all sessions
    pub total_usage: TokenUsage,
    
    /// Average tokens per session
    pub average_tokens_per_session: u32,
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[tokio::test]
    async fn test_token_monitor_basic() {
        let config = SummarizationConfig::default();
        let monitor = TokenMonitor::new(config);
        
        let usage = TokenUsage {
            input_tokens: 1000,
            output_tokens: 500,
            total_tokens: 1500,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        monitor.update_usage("test-session", usage).await.unwrap();
        
        let retrieved_usage = monitor.get_usage("test-session").await;
        assert_eq!(retrieved_usage.total_tokens, 1500);
        assert_eq!(retrieved_usage.input_tokens, 1000);
        assert_eq!(retrieved_usage.output_tokens, 500);
    }
    
    #[tokio::test]
    async fn test_cumulative_usage() {
        let config = SummarizationConfig::default();
        let monitor = TokenMonitor::new(config);
        
        let usage1 = TokenUsage {
            input_tokens: 1000,
            output_tokens: 500,
            total_tokens: 1500,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };

        let usage2 = TokenUsage {
            input_tokens: 800,
            output_tokens: 300,
            total_tokens: 1100,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        monitor.update_usage("test-session", usage1).await.unwrap();
        monitor.update_usage("test-session", usage2).await.unwrap();
        
        let total_usage = monitor.get_usage("test-session").await;
        assert_eq!(total_usage.total_tokens, 2600);
        assert_eq!(total_usage.input_tokens, 1800);
        assert_eq!(total_usage.output_tokens, 800);
    }
    
    #[tokio::test]
    async fn test_metrics_calculation() {
        let config = SummarizationConfig::default();
        let monitor = TokenMonitor::new(config);
        
        let usage = TokenUsage {
            input_tokens: 7000,
            output_tokens: 1000,
            total_tokens: 8000,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        monitor.update_usage("test-session", usage).await.unwrap();
        
        let metrics = monitor.get_metrics("test-session", "gpt-4", 10).await.unwrap();
        assert_eq!(metrics.total_tokens, 8000);
        assert_eq!(metrics.message_count, 10);
        assert_eq!(metrics.context_window, 8192);
        assert!(metrics.usage_percentage > 0.9); // Should be close to 98%
        assert!(metrics.needs_summarization); // Should trigger at 95%
    }
    
    #[tokio::test]
    async fn test_context_window_lookup() {
        let config = SummarizationConfig::default();
        let monitor = TokenMonitor::new(config);
        
        assert_eq!(monitor.get_context_window("gpt-4"), 8192);
        assert_eq!(monitor.get_context_window("gpt-4-turbo"), 128000);
        assert_eq!(monitor.get_context_window("claude-3-5-sonnet-20241022"), 200000);
        assert_eq!(monitor.get_context_window("unknown-model"), 4096);
    }
    
    #[tokio::test]
    async fn test_needs_summarization() {
        let config = SummarizationConfig {
            enabled: true,
            token_threshold: 5000,
            message_threshold: 20,
            context_window_percentage: 0.95,
            ..Default::default()
        };
        let monitor = TokenMonitor::new(config);
        
        // Test token threshold
        let usage = TokenUsage {
            input_tokens: 4000,
            output_tokens: 2000,
            total_tokens: 6000,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        monitor.update_usage("test-session", usage).await.unwrap();
        
        let needs = monitor.needs_summarization("test-session", "gpt-4", 5).await.unwrap();
        assert!(needs); // Should exceed token threshold
        
        // Test message threshold
        monitor.clear_session("test-session").await.unwrap();
        let small_usage = TokenUsage {
            input_tokens: 100,
            output_tokens: 50,
            total_tokens: 150,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        monitor.update_usage("test-session", small_usage).await.unwrap();
        
        let needs = monitor.needs_summarization("test-session", "gpt-4", 25).await.unwrap();
        assert!(needs); // Should exceed message threshold
    }
    
    #[tokio::test]
    async fn test_usage_stats() {
        let config = SummarizationConfig::default();
        let monitor = TokenMonitor::new(config);
        
        let usage1 = TokenUsage {
            input_tokens: 1000,
            output_tokens: 500,
            total_tokens: 1500,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        let usage2 = TokenUsage {
            input_tokens: 2000,
            output_tokens: 1000,
            total_tokens: 3000,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        monitor.update_usage("session1", usage1).await.unwrap();
        monitor.update_usage("session2", usage2).await.unwrap();
        
        let stats = monitor.get_usage_stats().await;
        assert_eq!(stats.session_count, 2);
        assert_eq!(stats.total_usage.total_tokens, 4500);
        assert_eq!(stats.average_tokens_per_session, 2250);
    }
}