oxirs-chat 0.3.2

use anyhow::Result;
use async_trait::async_trait;
use futures_util::StreamExt;
use oxirs_chat::llm::config::ProviderConfig;
use oxirs_chat::llm::providers::LLMProvider;
use oxirs_chat::llm::types::{LLMRequest, LLMResponse, LLMResponseChunk, LLMResponseStream, Usage};
use oxirs_chat::llm::LLMConfig;
use oxirs_chat::*;
use oxirs_core::ConcreteStore;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tempfile::TempDir;
use uuid::Uuid;

/// Mock LLM provider for testing
#[derive(Clone)]
struct MockLLMProvider {
    name: String,
}

impl MockLLMProvider {
    fn new(name: String) -> Self {
        Self { name }
    }
}

#[async_trait]
impl LLMProvider for MockLLMProvider {
    async fn generate(&self, model: &str, request: &LLMRequest) -> Result<LLMResponse> {
        // Simulate processing time
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;

        // Generate a mock response based on the input
        let response_content = format!(
            "Mock response to: {}",
            request
                .messages
                .first()
                .map(|m| m.content.as_str())
                .unwrap_or("")
        );

        Ok(LLMResponse {
            content: response_content,
            model_used: model.to_string(),
            provider_used: self.name.clone(),
            usage: Usage {
                prompt_tokens: 50,
                completion_tokens: 100,
                total_tokens: 150,
                cost: 0.001,
            },
            latency: std::time::Duration::from_millis(100),
            quality_score: Some(0.85),
            metadata: HashMap::new(),
        })
    }

    async fn generate_stream(
        &self,
        model: &str,
        request: &LLMRequest,
    ) -> Result<LLMResponseStream> {
        let started_at = Instant::now();

        // Build the same canned response used by `generate`, then split it
        // into four deterministic chunks so callers can verify accumulation.
        let full_content = format!(
            "Mock response to: {}",
            request
                .messages
                .first()
                .map(|m| m.content.as_str())
                .unwrap_or("")
        );

        // Split into roughly equal chunks; the last one carries `is_final`.
        let chunk_count = 4usize;
        let chars: Vec<char> = full_content.chars().collect();
        // Guard against empty content: at least one chunk of size 1.
        let chunk_size = ((chars.len() + chunk_count - 1) / chunk_count).max(1);

        let model_name = model.to_string();
        let provider_name = self.name.clone();

        // Collect into a Vec so we know the total before building is_final flags.
        let raw_chunks: Vec<Vec<char>> = chars.chunks(chunk_size).map(|s| s.to_vec()).collect();
        let total_chunks = raw_chunks.len();

        let chunks: Vec<Result<LLMResponseChunk>> = raw_chunks
            .into_iter()
            .enumerate()
            .map(|(idx, slice)| {
                let content: String = slice.into_iter().collect();
                let is_final = idx == total_chunks.saturating_sub(1);
                Ok(LLMResponseChunk {
                    content,
                    is_final,
                    chunk_index: idx,
                    model_used: model_name.clone(),
                    provider_used: provider_name.clone(),
                    latency: started_at.elapsed(),
                    metadata: HashMap::new(),
                })
            })
            .collect();

        let stream = futures_util::stream::iter(chunks);

        Ok(LLMResponseStream {
            stream: Box::pin(stream),
            model_used: model.to_string(),
            provider_used: self.name.clone(),
            started_at,
        })
    }

    fn get_available_models(&self) -> Vec<String> {
        vec!["mock-model".to_string()]
    }

    fn supports_streaming(&self) -> bool {
        true
    }

    fn get_provider_name(&self) -> &str {
        &self.name
    }

    fn estimate_cost(&self, _model: &str, _input_tokens: usize, _output_tokens: usize) -> f64 {
        0.001 // Mock cost
    }
}

/// Create a test-friendly LLM config with mock provider
#[allow(dead_code)]
fn create_test_llm_config() -> LLMConfig {
    let mut providers = HashMap::new();
    providers.insert(
        "mock".to_string(),
        ProviderConfig {
            enabled: true,
            api_key: Some("test-key".to_string()),
            base_url: Some("http://localhost:8080".to_string()),
            models: vec![oxirs_chat::llm::config::ModelConfig {
                name: "mock-model".to_string(),
                max_tokens: 1000,
                cost_per_token: 0.00001,
                capabilities: vec!["test".to_string()],
                use_cases: vec!["testing".to_string()],
            }],
            timeout: std::time::Duration::from_secs(30),
            max_retries: 3,
        },
    );

    LLMConfig {
        providers,
        routing: oxirs_chat::llm::config::RoutingConfig::default(),
        fallback: oxirs_chat::llm::config::FallbackConfig::default(),
        rate_limits: oxirs_chat::llm::config::RateLimitConfig::default(),
        circuit_breaker: oxirs_chat::llm::config::CircuitBreakerConfig::default(),
    }
}

/// Test LLM Manager that bypasses the provider initialization
pub struct TestLLMManager {
    mock_provider: MockLLMProvider,
}

impl Default for TestLLMManager {
    fn default() -> Self {
        Self::new()
    }
}

impl TestLLMManager {
    pub fn new() -> Self {
        Self {
            mock_provider: MockLLMProvider::new("mock".to_string()),
        }
    }

    pub async fn generate_response(&self, request: LLMRequest) -> Result<LLMResponse> {
        self.mock_provider.generate("mock-model", &request).await
    }
}

#[tokio::test]
async fn test_session_creation_and_persistence() {
    // Create a temporary directory for testing
    let temp_dir = TempDir::new().expect("Failed to create temp directory");
    let persistence_path = temp_dir.path().join("test_sessions");

    // Create a store for testing
    let store = Arc::new(ConcreteStore::new().expect("Failed to create store"));

    // Create chat manager with persistence
    let manager = ChatManager::with_persistence(store.clone(), &persistence_path)
        .await
        .expect("Failed to create chat manager");

    // Test session creation
    let session_id = "test_session_001".to_string();
    let session = manager
        .get_or_create_session(session_id.clone())
        .await
        .expect("Failed to create session");

    // Verify session exists
    assert_eq!(session.lock().await.id, session_id);

    // Add a message to the session
    {
        let mut session_guard = session.lock().await;
        // Create a test LLM manager
        let test_llm_manager = TestLLMManager::new();

        // For testing, we'll simulate the process_message functionality
        // Since we can't use the real process_message without modifying the Session struct
        // we'll manually create the expected behavior

        // Create user message
        let user_msg = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::User,
            content: oxirs_chat::messages::MessageContent::Text(
                "Hello, this is a test message".to_string(),
            ),
            timestamp: chrono::Utc::now(),
            metadata: None,
            thread_id: None,
            parent_message_id: None,
            token_count: Some(25),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };

        session_guard.messages.push(user_msg.clone());

        // Generate LLM response
        let llm_request = LLMRequest {
            messages: vec![oxirs_chat::llm::types::ChatMessage {
                role: oxirs_chat::llm::types::ChatRole::User,
                content: "Hello, this is a test message".to_string(),
                metadata: None,
            }],
            system_prompt: Some("You are a helpful assistant.".to_string()),
            max_tokens: Some(1000),
            temperature: 0.7,
            use_case: oxirs_chat::llm::types::UseCase::Conversation,
            priority: oxirs_chat::llm::types::Priority::Normal,
            timeout: Some(std::time::Duration::from_secs(30)),
        };

        let llm_response = test_llm_manager
            .generate_response(llm_request)
            .await
            .expect("Failed to generate LLM response");

        // Create assistant message
        let response = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::Assistant,
            content: oxirs_chat::messages::MessageContent::Text(llm_response.content),
            timestamp: chrono::Utc::now(),
            metadata: Some(oxirs_chat::messages::MessageMetadata {
                source: Some("test-llm".to_string()),
                confidence: Some(0.85),
                processing_time_ms: Some(100),
                model_used: Some("mock-model".to_string()),
                temperature: Some(0.7),
                max_tokens: Some(1000),
                custom_fields: std::collections::HashMap::new(),
            }),
            thread_id: None,
            parent_message_id: Some(user_msg.id),
            token_count: Some(100),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };

        session_guard.messages.push(response.clone());

        assert_eq!(response.role, oxirs_chat::messages::MessageRole::Assistant);
        assert!(!response.content.to_string().is_empty());
    }

    // Test session statistics
    let stats = manager.get_session_stats().await;
    assert_eq!(stats.total_sessions, 1);
    assert_eq!(stats.active_sessions, 1);

    // Test detailed metrics
    let metrics = manager.get_detailed_metrics().await;
    assert_eq!(metrics.total_sessions, 1);
    assert_eq!(metrics.active_sessions, 1);
    assert!(metrics.total_messages > 0);

    println!("✅ Session persistence test passed!");
}

#[tokio::test]
async fn test_session_backup_and_restore() {
    // Create temporary directories
    let temp_dir = TempDir::new().expect("Failed to create temp directory");
    let persistence_path = temp_dir.path().join("test_sessions");
    let backup_path = temp_dir.path().join("test_backup");

    // Create store and manager
    let store = Arc::new(ConcreteStore::new().expect("Failed to create store"));
    let manager = ChatManager::with_persistence(store.clone(), &persistence_path)
        .await
        .expect("Failed to create chat manager");

    // Create a test session with some data
    let session_id = "backup_test_session".to_string();
    let session = manager
        .get_or_create_session(session_id.clone())
        .await
        .expect("Failed to create session");

    // Add some messages to make the session interesting
    {
        let mut session_guard = session.lock().await;
        // Create a test LLM manager
        let test_llm_manager = TestLLMManager::new();

        // Add first message
        let user_msg1 = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::User,
            content: oxirs_chat::messages::MessageContent::Text("First test message".to_string()),
            timestamp: chrono::Utc::now(),
            metadata: None,
            thread_id: None,
            parent_message_id: None,
            token_count: Some(20),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        session_guard.messages.push(user_msg1.clone());

        let llm_request1 = LLMRequest {
            messages: vec![oxirs_chat::llm::types::ChatMessage {
                role: oxirs_chat::llm::types::ChatRole::User,
                content: "First test message".to_string(),
                metadata: None,
            }],
            system_prompt: Some("You are a helpful assistant.".to_string()),
            max_tokens: Some(1000),
            temperature: 0.7,
            use_case: oxirs_chat::llm::types::UseCase::Conversation,
            priority: oxirs_chat::llm::types::Priority::Normal,
            timeout: Some(std::time::Duration::from_secs(30)),
        };

        let llm_response1 = test_llm_manager
            .generate_response(llm_request1)
            .await
            .expect("Failed to generate first LLM response");

        let response1 = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::Assistant,
            content: oxirs_chat::messages::MessageContent::Text(llm_response1.content),
            timestamp: chrono::Utc::now(),
            metadata: Some(oxirs_chat::messages::MessageMetadata {
                source: Some("test-llm".to_string()),
                confidence: Some(0.85),
                processing_time_ms: Some(100),
                model_used: Some("mock-model".to_string()),
                temperature: Some(0.7),
                max_tokens: Some(1000),
                custom_fields: std::collections::HashMap::new(),
            }),
            thread_id: None,
            parent_message_id: Some(user_msg1.id),
            token_count: Some(100),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        session_guard.messages.push(response1);

        // Add second message
        let user_msg2 = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::User,
            content: oxirs_chat::messages::MessageContent::Text("Second test message".to_string()),
            timestamp: chrono::Utc::now(),
            metadata: None,
            thread_id: None,
            parent_message_id: None,
            token_count: Some(21),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        session_guard.messages.push(user_msg2.clone());

        let llm_request2 = LLMRequest {
            messages: vec![oxirs_chat::llm::types::ChatMessage {
                role: oxirs_chat::llm::types::ChatRole::User,
                content: "Second test message".to_string(),
                metadata: None,
            }],
            system_prompt: Some("You are a helpful assistant.".to_string()),
            max_tokens: Some(1000),
            temperature: 0.7,
            use_case: oxirs_chat::llm::types::UseCase::Conversation,
            priority: oxirs_chat::llm::types::Priority::Normal,
            timeout: Some(std::time::Duration::from_secs(30)),
        };

        let llm_response2 = test_llm_manager
            .generate_response(llm_request2)
            .await
            .expect("Failed to generate second LLM response");

        let response2 = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::Assistant,
            content: oxirs_chat::messages::MessageContent::Text(llm_response2.content),
            timestamp: chrono::Utc::now(),
            metadata: Some(oxirs_chat::messages::MessageMetadata {
                source: Some("test-llm".to_string()),
                confidence: Some(0.85),
                processing_time_ms: Some(100),
                model_used: Some("mock-model".to_string()),
                temperature: Some(0.7),
                max_tokens: Some(1000),
                custom_fields: std::collections::HashMap::new(),
            }),
            thread_id: None,
            parent_message_id: Some(user_msg2.id),
            token_count: Some(100),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        session_guard.messages.push(response2);
    }

    // Test backup
    let backup_report = manager
        .backup_sessions(&backup_path)
        .await
        .expect("Failed to backup sessions");

    assert_eq!(backup_report.successful_backups, 1);
    assert_eq!(backup_report.failed_backups, 0);

    // Clear sessions (simulate fresh start)
    manager
        .remove_session(&session_id)
        .await
        .expect("Failed to remove session");

    // Test restore
    let restore_report = manager
        .restore_sessions(&backup_path)
        .await
        .expect("Failed to restore sessions");

    assert_eq!(restore_report.sessions_restored, 1);
    // Note: failed_restorations field not available in RestoreReport

    // Verify the restored session
    let restored_session = manager
        .get_session(&session_id)
        .await
        .expect("Restored session not found");

    let restored_session_arc = restored_session.expect("Restored session should exist");
    let restored_session_guard = restored_session_arc.lock().await;
    assert_eq!(restored_session_guard.id, session_id);
    assert!(restored_session_guard.messages.len() >= 2); // Should have the messages we added

    println!("✅ Session backup and restore test passed!");
}

#[tokio::test]
async fn test_session_expiration_and_cleanup() {
    // Create temporary directory
    let temp_dir = TempDir::new().expect("Failed to create temp directory");
    let persistence_path = temp_dir.path().join("test_sessions");

    // Create store and manager
    let store = Arc::new(ConcreteStore::new().expect("Failed to create store"));
    let manager = ChatManager::with_persistence(store.clone(), &persistence_path)
        .await
        .expect("Failed to create chat manager");

    // Create a session with short timeout
    let session_id = "expiration_test_session".to_string();
    let session = manager
        .get_or_create_session(session_id.clone())
        .await
        .expect("Failed to create session");

    // Manually expire the session by setting old timestamp
    {
        let mut session_guard = session.lock().await;
        session_guard.last_activity = chrono::Utc::now() - chrono::Duration::hours(2);
    }

    // Test cleanup
    let cleaned_count = manager
        .cleanup_expired_sessions()
        .await
        .expect("Failed to cleanup expired sessions");

    assert_eq!(cleaned_count, 1);

    // Verify session was removed
    let session_result = manager.get_session(&session_id).await;
    assert!(session_result.expect("Should get result").is_none());

    println!("✅ Session expiration and cleanup test passed!");
}

/// Verify that `MockLLMProvider::generate_stream` yields chunks whose
/// accumulated text equals the non-streaming response, and that the
/// assembled assistant message can be stored in and retrieved from the
/// session layer exactly like a non-streamed message.
#[tokio::test]
async fn test_streaming_accumulation_and_session_persistence() {
    // ---- Setup -------------------------------------------------------
    let temp_dir = TempDir::new().expect("Failed to create temp directory");
    let persistence_path = temp_dir.path().join("stream_test_sessions");

    let store = Arc::new(ConcreteStore::new().expect("Failed to create store"));
    let manager = ChatManager::with_persistence(store.clone(), &persistence_path)
        .await
        .expect("Failed to create chat manager");

    let session_id = "stream_session_001".to_string();
    let session = manager
        .get_or_create_session(session_id.clone())
        .await
        .expect("Failed to create session");

    // ---- Stream through MockLLMProvider --------------------------------
    let provider = MockLLMProvider::new("mock".to_string());
    assert!(
        provider.supports_streaming(),
        "MockLLMProvider must advertise streaming support"
    );

    let input_text = "Hello from the streaming test";

    let request = LLMRequest {
        messages: vec![oxirs_chat::llm::types::ChatMessage {
            role: oxirs_chat::llm::types::ChatRole::User,
            content: input_text.to_string(),
            metadata: None,
        }],
        system_prompt: Some("You are a helpful assistant.".to_string()),
        max_tokens: Some(1000),
        temperature: 0.7,
        use_case: oxirs_chat::llm::types::UseCase::Conversation,
        priority: oxirs_chat::llm::types::Priority::Normal,
        timeout: Some(std::time::Duration::from_secs(30)),
    };

    // Obtain the stream and collect all chunks.
    let mut stream_wrapper = provider
        .generate_stream("mock-model", &request)
        .await
        .expect("generate_stream must not fail on MockLLMProvider");

    let mut accumulated = String::new();
    let mut chunk_count = 0usize;
    let mut saw_final = false;

    while let Some(chunk_result) = stream_wrapper.stream.next().await {
        let chunk = chunk_result.expect("Stream chunk must not be an error");
        accumulated.push_str(&chunk.content);
        chunk_count += 1;
        if chunk.is_final {
            saw_final = true;
        }
        // chunk_index must be strictly monotone within a single stream.
        assert_eq!(
            chunk.chunk_index,
            chunk_count - 1,
            "chunk_index must equal the iteration counter"
        );
        assert_eq!(chunk.model_used, "mock-model");
        assert_eq!(chunk.provider_used, "mock");
    }

    // Must have emitted at least one chunk and signalled completion.
    assert!(chunk_count > 0, "Stream must yield at least one chunk");
    assert!(saw_final, "Last chunk must have is_final == true");

    // The accumulated text must be identical to what non-streaming `generate`
    // would return for the same request.
    let expected_content = format!("Mock response to: {}", input_text);
    assert_eq!(
        accumulated, expected_content,
        "Accumulated streaming text must equal the non-streaming response"
    );

    // ---- Persist the assembled assistant message ----------------------
    {
        let mut guard = session.lock().await;

        let user_msg = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::User,
            content: oxirs_chat::messages::MessageContent::Text(input_text.to_string()),
            timestamp: chrono::Utc::now(),
            metadata: None,
            thread_id: None,
            parent_message_id: None,
            token_count: Some(20),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        guard.messages.push(user_msg.clone());

        let assistant_msg = oxirs_chat::messages::Message {
            id: format!("msg_{}", Uuid::new_v4()),
            role: oxirs_chat::messages::MessageRole::Assistant,
            content: oxirs_chat::messages::MessageContent::Text(accumulated.clone()),
            timestamp: chrono::Utc::now(),
            metadata: Some(oxirs_chat::messages::MessageMetadata {
                source: Some("mock-stream".to_string()),
                confidence: Some(0.85),
                processing_time_ms: Some(50),
                model_used: Some("mock-model".to_string()),
                temperature: Some(0.7),
                max_tokens: Some(1000),
                custom_fields: std::collections::HashMap::new(),
            }),
            thread_id: None,
            parent_message_id: Some(user_msg.id),
            token_count: Some(chunk_count),
            reactions: Vec::new(),
            attachments: Vec::new(),
            rich_elements: Vec::new(),
        };
        guard.messages.push(assistant_msg);
    }

    // ---- Verify the session layer recorded the message -----------------
    let stats = manager.get_session_stats().await;
    assert_eq!(stats.total_sessions, 1);
    assert_eq!(stats.active_sessions, 1);

    let metrics = manager.get_detailed_metrics().await;
    assert!(
        metrics.total_messages >= 2,
        "Session must contain at least user + assistant messages"
    );

    // Re-read the session and confirm accumulated content was stored.
    let stored = manager
        .get_session(&session_id)
        .await
        .expect("Session lookup must not error")
        .expect("Session must still exist after writing messages");

    let stored_guard = stored.lock().await;
    let assistant_stored = stored_guard
        .messages
        .iter()
        .find(|m| m.role == oxirs_chat::messages::MessageRole::Assistant)
        .expect("Assistant message must be present in persisted session");

    assert_eq!(
        assistant_stored.content.to_string(),
        expected_content,
        "Persisted assistant content must match the accumulated stream text"
    );

    println!("✅ Streaming accumulation and session persistence test passed!");
}