//! RAG Pipeline Integration Tests
//!
//! Tests the complete RAG pipeline end-to-end with mock components.

use std::sync::Arc;

use aifastdb_core::rag::{
    ChatMessage, ChatRole, DefaultContextAssembler, InMemorySessionManager, RagConfig,
    RagPipeline, RagResponse, RagResult, RagTokenUsage, RetrievalResult,
};
use aifastdb_core::rag::traits::{Generator, Retriever, SessionManager};

// ============================================================================
// Mock Components
// ============================================================================

/// Mock retriever that returns pre-configured documents.
struct MockRetriever {
    documents: Vec<RetrievalResult>,
}

impl MockRetriever {
    fn new(docs: Vec<(&str, &str, f64)>) -> Self {
        Self {
            documents: docs
                .into_iter()
                .map(|(id, content, score)| RetrievalResult::new(id, content, score))
                .collect(),
        }
    }
}

impl Retriever for MockRetriever {
    fn retrieve(
        &self,
        _query: &str,
        _embedding: Option<&[f32]>,
        config: &RagConfig,
    ) -> RagResult<Vec<RetrievalResult>> {
        Ok(self
            .documents
            .iter()
            .take(config.top_k)
            .cloned()
            .collect())
    }
}

/// Mock generator that produces deterministic answers based on the last user message.
struct EchoGenerator;

impl Generator for EchoGenerator {
    fn generate(
        &self,
        messages: &[ChatMessage],
        _config: &RagConfig,
    ) -> RagResult<RagResponse> {
        let user_msg = messages
            .iter()
            .rev()
            .find(|m| m.role == ChatRole::User)
            .map(|m| m.content.as_str())
            .unwrap_or("unknown");

        // Count context documents in system message (format: "[1]", "[2]", etc.)
        let doc_count = messages
            .iter()
            .find(|m| m.role == ChatRole::System)
            .map(|m| {
                // Count lines matching pattern [N] (score: X.XX)
                m.content
                    .lines()
                    .filter(|line| {
                        let trimmed = line.trim();
                        trimmed.starts_with('[')
                            && trimmed.contains("(score:")
                    })
                    .count()
            })
            .unwrap_or(0);

        Ok(RagResponse {
            answer: format!(
                "Based on {} documents, here is the answer to: {}",
                doc_count, user_msg
            ),
            sources: Vec::new(),
            session_id: None,
            usage: Some(RagTokenUsage {
                prompt_tokens: 200,
                completion_tokens: 50,
                total_tokens: 250,
                context_tokens: 0,
                cost: None,
                reasoning_tokens: None,
            }),
            retrieval_count: 0,
            latency_ms: None,
            metadata: Default::default(),
        })
    }
}

// ============================================================================
// Integration Tests
// ============================================================================

#[test]
fn test_full_rag_pipeline_one_shot() {
    let retriever = MockRetriever::new(vec![
        ("doc-1", "Rust is a systems programming language focused on safety.", 0.95),
        ("doc-2", "Rust has zero-cost abstractions and memory safety.", 0.88),
        ("doc-3", "Cargo is Rust's package manager and build tool.", 0.75),
    ]);

    let pipeline = RagPipeline::builder()
        .retriever(retriever)
        .generator(EchoGenerator)
        .config(RagConfig {
            top_k: 3,
            include_sources: true,
            ..Default::default()
        })
        .build()
        .expect("Pipeline should build successfully");

    let response = pipeline.query("What is Rust?").unwrap();

    // Verify response
    assert!(response.answer.contains("What is Rust?"));
    assert!(response.answer.contains("3 documents"));
    assert_eq!(response.retrieval_count, 3);
    assert_eq!(response.sources.len(), 3);
    assert!(response.latency_ms.is_some());
    assert!(response.session_id.is_none()); // one-shot has no session

    // Verify sources
    assert_eq!(response.sources[0].id, "doc-1");
    assert!(response.sources[0].score > 0.9);

    // Verify token usage
    let usage = response.usage.unwrap();
    assert_eq!(usage.prompt_tokens, 200);
    assert_eq!(usage.completion_tokens, 50);
}

#[test]
fn test_full_rag_pipeline_multi_turn_conversation() {
    let retriever = MockRetriever::new(vec![
        ("doc-a", "AiFastDB is a high-performance vector database.", 0.92),
        ("doc-b", "AiFastDB supports hybrid search combining vector and text.", 0.85),
    ]);

    let sessions: Arc<dyn SessionManager> = Arc::new(InMemorySessionManager::default());

    let pipeline = RagPipeline::builder()
        .retriever(retriever)
        .generator(EchoGenerator)
        .session_manager(Arc::clone(&sessions))
        .config(RagConfig {
            top_k: 2,
            include_sources: true,
            max_history_turns: 5,
            ..Default::default()
        })
        .build()
        .unwrap();

    // Turn 1
    let r1 = pipeline.chat("session-001", "What is AiFastDB?").unwrap();
    assert!(r1.answer.contains("What is AiFastDB?"));
    assert_eq!(r1.session_id.as_deref(), Some("session-001"));
    assert_eq!(r1.retrieval_count, 2);

    // Turn 2 (same session)
    let r2 = pipeline
        .chat("session-001", "What search modes does it support?")
        .unwrap();
    assert!(r2.answer.contains("What search modes does it support?"));

    // Verify session history (4 messages: 2 user + 2 assistant)
    let history = sessions.get_history("session-001").unwrap();
    assert_eq!(history.len(), 4);
    assert_eq!(history[0].role, ChatRole::User);
    assert_eq!(history[0].content, "What is AiFastDB?");
    assert_eq!(history[1].role, ChatRole::Assistant);
    assert_eq!(history[2].role, ChatRole::User);
    assert_eq!(history[2].content, "What search modes does it support?");
    assert_eq!(history[3].role, ChatRole::Assistant);
}

#[test]
fn test_session_isolation() {
    let retriever = MockRetriever::new(vec![
        ("doc-x", "Some content about X.", 0.9),
    ]);

    let sessions: Arc<dyn SessionManager> = Arc::new(InMemorySessionManager::default());

    let pipeline = RagPipeline::builder()
        .retriever(retriever)
        .generator(EchoGenerator)
        .session_manager(Arc::clone(&sessions))
        .config(RagConfig::default())
        .build()
        .unwrap();

    // Two separate sessions
    pipeline.chat("alice", "Hello from Alice").unwrap();
    pipeline.chat("bob", "Hello from Bob").unwrap();
    pipeline.chat("alice", "Alice's second message").unwrap();

    assert_eq!(sessions.session_count(), 2);

    let alice_history = sessions.get_history("alice").unwrap();
    let bob_history = sessions.get_history("bob").unwrap();

    assert_eq!(alice_history.len(), 4); // 2 user + 2 assistant
    assert_eq!(bob_history.len(), 2); // 1 user + 1 assistant

    assert_eq!(alice_history[0].content, "Hello from Alice");
    assert_eq!(alice_history[2].content, "Alice's second message");
    assert_eq!(bob_history[0].content, "Hello from Bob");
}

#[test]
fn test_pipeline_without_sources() {
    let retriever = MockRetriever::new(vec![
        ("doc-1", "Content 1", 0.9),
        ("doc-2", "Content 2", 0.8),
    ]);

    let pipeline = RagPipeline::builder()
        .retriever(retriever)
        .generator(EchoGenerator)
        .config(RagConfig {
            include_sources: false,
            ..Default::default()
        })
        .build()
        .unwrap();

    let response = pipeline.query("test").unwrap();
    assert!(response.sources.is_empty());
    assert_eq!(response.retrieval_count, 2);
}

#[test]
fn test_pipeline_builder_validation() {
    // Missing retriever → error
    let result = RagPipeline::builder().generator(EchoGenerator).build();
    assert!(result.is_err());
    assert!(result.unwrap_err().to_string().contains("Retriever"));

    // Missing generator → error
    let result = RagPipeline::builder()
        .retriever(MockRetriever::new(vec![]))
        .build();
    assert!(result.is_err());
    assert!(result.unwrap_err().to_string().contains("Generator"));

    // Both provided → ok
    let result = RagPipeline::builder()
        .retriever(MockRetriever::new(vec![]))
        .generator(EchoGenerator)
        .build();
    assert!(result.is_ok());
}

#[test]
fn test_config_defaults() {
    let config = RagConfig::default();
    assert!(config.top_k > 0);
    assert!(config.max_context_tokens > 0);
    assert!(config.max_history_turns > 0);
    assert!(config.include_sources);
    assert!(config.temperature >= 0.0 && config.temperature <= 2.0);
}

#[test]
fn test_context_assembly_token_budgeting() {
    // Create a retriever with long documents
    let long_doc = "a ".repeat(5000); // ~2500 tokens
    let retriever = MockRetriever::new(vec![
        ("long-1", &long_doc, 0.9),
        ("long-2", &long_doc, 0.8),
    ]);

    let pipeline = RagPipeline::builder()
        .retriever(retriever)
        .generator(EchoGenerator)
        .config(RagConfig {
            max_context_tokens: 1000, // Very tight budget
            top_k: 2,
            ..Default::default()
        })
        .build()
        .unwrap();

    // Should not panic even with very long docs
    let response = pipeline.query("Summarize the docs").unwrap();
    assert!(!response.answer.is_empty());
}

#[test]
fn test_session_management_operations() {
    let sessions = InMemorySessionManager::default();

    // Create sessions
    let s1 = sessions.create_session().unwrap();
    let s2 = sessions.create_session_with_id("custom-id").unwrap();

    assert_eq!(sessions.session_count(), 2);
    assert_eq!(s2.session_id, "custom-id");

    // Add messages
    sessions
        .push_message(&s1.session_id, ChatMessage::user("Hello"))
        .unwrap();
    sessions
        .push_message(&s1.session_id, ChatMessage::assistant("Hi there"))
        .unwrap();

    // Get history
    let history = sessions.get_history(&s1.session_id).unwrap();
    assert_eq!(history.len(), 2);

    // Get recent history
    let recent = sessions.get_recent_history(&s1.session_id, 1).unwrap();
    assert_eq!(recent.len(), 2); // 1 turn = 2 messages (user + assistant)

    // List sessions
    let session_list = sessions.list_sessions();
    assert_eq!(session_list.len(), 2);

    // Delete session
    sessions.delete_session(&s1.session_id).unwrap();
    assert_eq!(sessions.session_count(), 1);

    // Deleted session should not be accessible
    assert!(sessions.get_history(&s1.session_id).is_err());
}
