coderlib 0.1.0

//! Session management for CoderLib
//!
//! This module handles creating, managing, and persisting conversation sessions
//! between users and AI agents.

use std::sync::Arc;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::core::error::{SessionError, SessionResult};
use crate::storage::{Storage, Message};

/// A conversation session between a user and AI agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    /// Unique session identifier
    pub id: String,
    
    /// Human-readable session title
    pub title: String,
    
    /// When the session was created
    pub created_at: DateTime<Utc>,
    
    /// When the session was last updated
    pub updated_at: DateTime<Utc>,
    
    /// AI model used for this session
    pub model: String,
    
    /// Total cost for this session (in USD)
    pub cost: f64,
    
    /// Token usage statistics
    pub token_usage: TokenUsage,
    
    /// Session metadata
    pub metadata: serde_json::Value,
}

/// Token usage statistics for a session
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TokenUsage {
    /// Total input tokens used
    pub input_tokens: u32,
    
    /// Total output tokens generated
    pub output_tokens: u32,
    
    /// Cache creation tokens (if applicable)
    pub cache_creation_tokens: u32,
    
    /// Cache read tokens (if applicable)
    pub cache_read_tokens: u32,
}

/// Session manager handles session lifecycle
pub struct SessionManager {
    storage: Arc<dyn Storage>,
}

impl SessionManager {
    /// Create a new session manager
    pub fn new(storage: Arc<dyn Storage>) -> Self {
        Self { storage }
    }
    
    /// Create a new session
    pub async fn create_session(&self, title: Option<String>) -> SessionResult<String> {
        let session_id = Uuid::new_v4().to_string();
        let now = Utc::now();
        
        let session = Session {
            id: session_id.clone(),
            title: title.unwrap_or_else(|| format!("Session {}", &session_id[..8])),
            created_at: now,
            updated_at: now,
            model: "unknown".to_string(), // Will be updated when first message is sent
            cost: 0.0,
            token_usage: TokenUsage::default(),
            metadata: serde_json::Value::Null,
        };
        
        self.storage.create_session(&session).await
            .map_err(|e| SessionError::CreationFailed(e.to_string()))?;
        
        Ok(session_id)
    }
    
    /// Get a session by ID
    pub async fn get_session(&self, session_id: &str) -> SessionResult<Option<Session>> {
        self.storage.get_session(session_id).await
            .map_err(|e| SessionError::NotFound(e.to_string()))
    }
    
    /// Update a session
    pub async fn update_session(&self, session: &Session) -> SessionResult<()> {
        self.storage.update_session(session).await
            .map_err(|e| SessionError::CreationFailed(e.to_string()))
    }
    
    /// Delete a session
    pub async fn delete_session(&self, session_id: &str) -> SessionResult<()> {
        self.storage.delete_session(session_id).await
            .map_err(|e| SessionError::NotFound(e.to_string()))
    }
    
    /// List all sessions
    pub async fn list_sessions(&self) -> SessionResult<Vec<Session>> {
        self.storage.list_sessions().await
            .map_err(|e| SessionError::NotFound(e.to_string()))
    }
    
    /// Get messages for a session
    pub async fn get_messages(&self, session_id: &str) -> SessionResult<Vec<Message>> {
        self.storage.get_messages(session_id).await
            .map_err(|e| SessionError::NotFound(e.to_string()))
    }
    
    /// Add a message to a session
    pub async fn add_message(&self, session_id: &str, message: &Message) -> SessionResult<String> {
        // Verify session exists
        if self.get_session(session_id).await?.is_none() {
            return Err(SessionError::NotFound(session_id.to_string()));
        }
        
        let message_id = self.storage.save_message(message).await
            .map_err(|e| SessionError::CreationFailed(e.to_string()))?;
        
        // Update session timestamp
        if let Some(mut session) = self.get_session(session_id).await? {
            session.updated_at = Utc::now();
            self.update_session(&session).await?;
        }
        
        Ok(message_id)
    }
    
    /// Update token usage for a session
    pub async fn update_token_usage(
        &self,
        session_id: &str,
        token_usage: &crate::core::TokenUsage,
    ) -> SessionResult<()> {
        if let Some(mut session) = self.get_session(session_id).await? {
            session.token_usage.input_tokens += token_usage.input_tokens;
            session.token_usage.output_tokens += token_usage.output_tokens;
            session.token_usage.cache_creation_tokens += token_usage.cache_creation_tokens;
            session.token_usage.cache_read_tokens += token_usage.cache_read_tokens;
            session.updated_at = Utc::now();
            
            self.update_session(&session).await?;
        } else {
            return Err(SessionError::NotFound(session_id.to_string()));
        }
        
        Ok(())
    }
    
    /// Update cost for a session
    pub async fn update_cost(&self, session_id: &str, additional_cost: f64) -> SessionResult<()> {
        if let Some(mut session) = self.get_session(session_id).await? {
            session.cost += additional_cost;
            session.updated_at = Utc::now();
            
            self.update_session(&session).await?;
        } else {
            return Err(SessionError::NotFound(session_id.to_string()));
        }
        
        Ok(())
    }
    
    /// Set the model for a session
    pub async fn set_model(&self, session_id: &str, model: &str) -> SessionResult<()> {
        if let Some(mut session) = self.get_session(session_id).await? {
            session.model = model.to_string();
            session.updated_at = Utc::now();
            
            self.update_session(&session).await?;
        } else {
            return Err(SessionError::NotFound(session_id.to_string()));
        }
        
        Ok(())
    }
    
    /// Update session title
    pub async fn update_title(&self, session_id: &str, title: &str) -> SessionResult<()> {
        if let Some(mut session) = self.get_session(session_id).await? {
            session.title = title.to_string();
            session.updated_at = Utc::now();
            
            self.update_session(&session).await?;
        } else {
            return Err(SessionError::NotFound(session_id.to_string()));
        }
        
        Ok(())
    }
    
    /// Get session statistics
    pub async fn get_statistics(&self, session_id: &str) -> SessionResult<SessionStatistics> {
        let session = self.get_session(session_id).await?
            .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
        
        let messages = self.get_messages(session_id).await?;
        
        let user_messages = messages.iter()
            .filter(|m| m.role == crate::storage::MessageRole::User)
            .count();
        
        let assistant_messages = messages.iter()
            .filter(|m| m.role == crate::storage::MessageRole::Assistant)
            .count();
        
        Ok(SessionStatistics {
            session_id: session_id.to_string(),
            total_messages: messages.len(),
            user_messages,
            assistant_messages,
            total_tokens: session.token_usage.input_tokens + session.token_usage.output_tokens,
            total_cost: session.cost,
            duration: Utc::now().signed_duration_since(session.created_at),
        })
    }
}

/// Statistics for a session
#[derive(Debug, Clone)]
pub struct SessionStatistics {
    pub session_id: String,
    pub total_messages: usize,
    pub user_messages: usize,
    pub assistant_messages: usize,
    pub total_tokens: u32,
    pub total_cost: f64,
    pub duration: chrono::Duration,
}

impl TokenUsage {
    /// Calculate total tokens
    pub fn total(&self) -> u32 {
        self.input_tokens + self.output_tokens + self.cache_creation_tokens + self.cache_read_tokens
    }
    
    /// Add token usage from another instance
    pub fn add(&mut self, other: &TokenUsage) {
        self.input_tokens += other.input_tokens;
        self.output_tokens += other.output_tokens;
        self.cache_creation_tokens += other.cache_creation_tokens;
        self.cache_read_tokens += other.cache_read_tokens;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    // Note: MockStorage will be implemented when we create the storage module

    // TODO: Implement test when MockStorage is available

    #[test]
    fn test_token_usage() {
        let mut usage = TokenUsage::default();
        assert_eq!(usage.total(), 0);
        
        usage.input_tokens = 100;
        usage.output_tokens = 50;
        assert_eq!(usage.total(), 150);
        
        let other = TokenUsage {
            input_tokens: 25,
            output_tokens: 25,
            cache_creation_tokens: 10,
            cache_read_tokens: 5,
        };
        
        usage.add(&other);
        assert_eq!(usage.input_tokens, 125);
        assert_eq!(usage.output_tokens, 75);
        assert_eq!(usage.total(), 215);
    }

    #[test]
    fn test_session_creation() {
        let session = Session {
            id: "test-id".to_string(),
            title: "Test Session".to_string(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
            model: "gpt-4".to_string(),
            cost: 0.0,
            token_usage: TokenUsage::default(),
            metadata: serde_json::Value::Null,
        };
        
        assert_eq!(session.id, "test-id");
        assert_eq!(session.title, "Test Session");
        assert_eq!(session.model, "gpt-4");
    }
}