//! Complete RAG (Retrieval-Augmented Generation) demonstration
//!
//! This example shows the full RTX RAG pipeline in action:
//! 1. Document indexing with different chunking strategies
//! 2. Dense retrieval with embedding generation
//! 3. Vector similarity search with filtering
//! 4. Integration with transformer generation
//! 5. Advanced features like hybrid search and re-ranking
//!
//! Run with: cargo run --example rag_complete_demo

use anyhow::Result;
use rtx_transformers::prelude::*;
use std::collections::HashMap;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize RTX Transformers
    rtx_transformers::init()?;

    println!("🚀 RTX RAG Complete Demonstration");
    println!("==================================\n");

    // Run the complete RAG demonstration
    run_rag_demo().await?;

    println!("\n✅ RAG demonstration completed successfully!");
    Ok(())
}

async fn run_rag_demo() -> Result<()> {
    // Step 1: Configure RAG Pipeline
    println!("📋 Step 1: Configuring RAG Pipeline");
    println!("-----------------------------------");

    let rag_config = RAGConfig {
        retriever_config: DenseRetrieverConfig {
            model_name: "sentence-transformers/all-MiniLM-L6-v2".to_string(),
            embedding_dim: 384,
            max_sequence_length: 512,
            batch_size: 16,
            device: Device::cuda(0).unwrap_or(Device::default()),
        },
        vector_db_config: VectorDbConfig {
            embedding_dim: 384,
            similarity_metric: SimilarityMetric::Cosine,
            index_type: IndexType::FlatL2,
        },
        chunking_strategy: ChunkingStrategy::FixedSize {
            size: 300,
            overlap: 75,
        },
        top_k: 5,
        rerank: true,
        query_expansion: false,
    };

    println!(
        "  ✓ Embedding dimension: {}",
        rag_config.vector_db_config.embedding_dim
    );
    println!(
        "  ✓ Similarity metric: {:?}",
        rag_config.vector_db_config.similarity_metric
    );
    println!("  ✓ Chunking strategy: {:?}", rag_config.chunking_strategy);
    println!("  ✓ Top-K retrieval: {}", rag_config.top_k);

    // Step 2: Create RAG Pipeline
    println!("\n🔧 Step 2: Creating RAG Pipeline");
    println!("--------------------------------");

    let mut pipeline = RAGPipeline::new(rag_config).await?;
    println!("  ✓ RAG pipeline initialized successfully");

    // Step 3: Prepare Knowledge Base
    println!("\n📚 Step 3: Preparing Knowledge Base");
    println!("-----------------------------------");

    let knowledge_base = create_sample_knowledge_base();
    println!(
        "  ✓ Created {} documents for indexing",
        knowledge_base.len()
    );

    for (i, doc) in knowledge_base.iter().enumerate() {
        println!(
            "    {}. {} ({})",
            i + 1,
            doc.metadata.get("title").unwrap_or(&"Untitled".to_string()),
            doc.metadata
                .get("category")
                .unwrap_or(&"Unknown".to_string())
        );
    }

    // Step 4: Index Documents
    println!("\n⚡ Step 4: Indexing Documents");
    println!("----------------------------");

    let start_time = std::time::Instant::now();
    let indexing_stats = pipeline.index_documents(&knowledge_base).await?;
    let indexing_duration = start_time.elapsed();

    println!(
        "  ✓ Documents processed: {}",
        indexing_stats.documents_processed
    );
    println!("  ✓ Chunks created: {}", indexing_stats.chunks_created);
    println!(
        "  ✓ Processing time: {:.2}ms",
        indexing_stats.processing_time_ms
    );
    println!(
        "  ✓ Avg chunks/doc: {:.1}",
        indexing_stats.avg_chunks_per_doc
    );
    println!("  ✓ Total indexing time: {:?}", indexing_duration);

    // Step 5: Demonstrate Retrieval
    println!("\n🔍 Step 5: Demonstrating Retrieval");
    println!("----------------------------------");

    let test_queries = vec![
        "What is machine learning and how does it work?",
        "Explain neural networks and deep learning",
        "How do transformers process natural language?",
        "What are the benefits of Rust programming language?",
        "Tell me about vector databases and similarity search",
    ];

    for (i, query) in test_queries.iter().enumerate() {
        println!("\nQuery {}: \"{}\"", i + 1, query);
        println!("{}", "-".repeat(50));

        let search_start = std::time::Instant::now();
        let search_results = pipeline.search(query, None).await?;
        let search_duration = search_start.elapsed();

        println!(
            "  Search time: {:?} ({:.2}ms)",
            search_duration, search_results.query_time_ms
        );
        println!("  Retrieved {} chunks:", search_results.chunks.len());

        for (j, result) in search_results.chunks.iter().take(3).enumerate() {
            println!(
                "    {}. [Score: {:.3}] {}...",
                j + 1,
                result.score,
                result.chunk.content.chars().take(80).collect::<String>()
            );
        }
    }

    // Step 6: Demonstrate Filtering
    println!("\n🎯 Step 6: Demonstrating Filtered Search");
    println!("----------------------------------------");

    let ai_filter = SearchFilter {
        metadata_filters: HashMap::from([("category".to_string(), "AI".to_string())]),
        min_score: Some(0.1),
        max_results: Some(3),
    };

    let filtered_query = "artificial intelligence systems";
    println!("Query: \"{}\" (AI category only)", filtered_query);

    let filtered_results = pipeline.search(filtered_query, Some(ai_filter)).await?;
    println!(
        "  Filtered results: {} chunks",
        filtered_results.chunks.len()
    );

    for (i, result) in filtered_results.chunks.iter().enumerate() {
        let category = result
            .chunk
            .metadata
            .get("category")
            .unwrap_or(&"Unknown".to_string());
        println!(
            "    {}. [{}] Score: {:.3} - {}...",
            i + 1,
            category,
            result.score,
            result.chunk.content.chars().take(60).collect::<String>()
        );
    }

    // Step 7: Demonstrate RAG Generation
    println!("\n🤖 Step 7: Demonstrating RAG Generation");
    println!("---------------------------------------");

    let generation_config = GenerationConfig {
        max_tokens: 200,
        temperature: 0.7,
        top_k: 50,
        top_p: 0.9,
        include_context: true,
        max_context_length: 1000,
    };

    let rag_generator =
        RAGGenerationPipeline::new(pipeline.config().clone(), generation_config).await?;

    // Index the same documents in the generation pipeline
    let mut rag_gen = rag_generator;
    rag_gen.index_documents(&knowledge_base).await?;

    let generation_query = "What is machine learning and why is it important?";
    println!("Generation Query: \"{}\"", generation_query);

    let generated_response = rag_gen.generate(generation_query, None).await?;

    println!("\nGenerated Response:");
    println!("{}", "=".repeat(60));
    println!("{}", generated_response.text);
    println!("{}", "=".repeat(60));

    println!(
        "\nContext Sources ({} chunks):",
        generated_response.context_chunks.len()
    );
    for (i, (chunk, score)) in generated_response
        .context_chunks
        .iter()
        .zip(generated_response.context_scores.iter())
        .enumerate()
    {
        let title = chunk
            .metadata
            .get("title")
            .unwrap_or(&"Unknown".to_string());
        println!(
            "  {}. [Score: {:.3}] {} - {}...",
            i + 1,
            score,
            title,
            chunk.content.chars().take(50).collect::<String>()
        );
    }

    // Step 8: Advanced Features Demo
    println!("\n🚀 Step 8: Advanced Features");
    println!("-----------------------------");

    // Query expansion
    println!("\n🔍 Query Expansion:");
    let query_expander = QueryExpander::new(QueryExpansionConfig {
        method: ExpansionMethod::Synonyms,
        max_expanded_queries: 3,
    })?;

    let original_query = "machine learning";
    let expanded_queries = query_expander.expand_query(original_query).await?;
    println!("  Original: \"{}\"", original_query);
    for (i, expanded) in expanded_queries.iter().enumerate() {
        println!("  Expanded {}: \"{}\"", i + 1, expanded);
    }

    // Context compression
    println!("\n🗜️  Context Compression:");
    let compressor = ContextCompressor::new(CompressionConfig {
        max_context_length: 500,
        compression_ratio: 0.6,
        preserve_query_relevance: true,
    })?;

    let search_results = pipeline.search("neural networks", None).await?;
    let chunks: Vec<DocumentChunk> = search_results
        .chunks
        .iter()
        .map(|r| r.chunk.clone())
        .collect();

    let compressed = compressor
        .compress_context(&chunks, "neural networks")
        .await?;
    println!("  Original chunks: {}", chunks.len());
    println!("  Compressed chunks: {}", compressed.len());
    println!(
        "  Compression ratio: {:.1}%",
        (compressed.len() as f32 / chunks.len() as f32) * 100.0
    );

    // Step 9: Performance Statistics
    println!("\n📊 Step 9: Performance Statistics");
    println!("----------------------------------");

    let final_stats = pipeline.get_stats().await?;
    println!("Pipeline Statistics:");
    for (key, value) in &final_stats {
        println!("  {}: {}", key, value);
    }

    // Memory usage estimation
    let chunk_count: usize = final_stats
        .get("total_chunks")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);
    let embedding_dim: usize = final_stats
        .get("embedding_dim")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);

    let estimated_memory_mb = (chunk_count * embedding_dim * 4) as f64 / (1024.0 * 1024.0);
    println!(
        "  Estimated embeddings memory: {:.2} MB",
        estimated_memory_mb
    );

    println!("\n✨ RAG System Features Demonstrated:");
    println!("  ✓ Document chunking with overlap");
    println!("  ✓ Dense embedding generation");
    println!("  ✓ Vector similarity search");
    println!("  ✓ Metadata filtering");
    println!("  ✓ RAG-based text generation");
    println!("  ✓ Query expansion");
    println!("  ✓ Context compression");
    println!("  ✓ Performance monitoring");

    Ok(())
}

fn create_sample_knowledge_base() -> Vec<Document> {
    vec![
        Document {
            id: "ml_fundamentals".to_string(),
            content: "Machine learning is a subset of artificial intelligence (AI) that provides systems the ability to automatically learn and improve from experience without being explicitly programmed. Machine learning focuses on the development of computer programs that can access data and use it to learn for themselves. The process of learning begins with observations or data, such as examples, direct experience, or instruction, in order to look for patterns in data and make better decisions in the future based on the examples that we provide.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Machine Learning Fundamentals".to_string()),
                ("category".to_string(), "AI".to_string()),
                ("difficulty".to_string(), "beginner".to_string()),
                ("topics".to_string(), "machine learning, AI, data science".to_string()),
            ]),
        },
        Document {
            id: "neural_networks_deep".to_string(),
            content: "Neural networks are computing systems vaguely inspired by the biological neural networks that constitute animal brains. Such systems learn to perform tasks by considering examples, generally without being programmed with task-specific rules. A neural network is based on a collection of connected units or nodes called artificial neurons, which loosely model the neurons in a biological brain. Deep learning is part of a broader family of machine learning methods based on artificial neural networks with representation learning. Learning can be supervised, semi-supervised or unsupervised.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Neural Networks and Deep Learning".to_string()),
                ("category".to_string(), "AI".to_string()),
                ("difficulty".to_string(), "intermediate".to_string()),
                ("topics".to_string(), "neural networks, deep learning, artificial neurons".to_string()),
            ]),
        },
        Document {
            id: "transformers_nlp".to_string(),
            content: "The Transformer is a deep learning model introduced in 2017 that is primarily used for natural language processing (NLP) tasks. Unlike recurrent neural networks (RNNs) that process tokens sequentially, the Transformer processes all tokens in a sequence simultaneously. This is made possible by the self-attention mechanism, which allows the model to weigh the importance of different words in a sentence when processing each word. The Transformer architecture has led to significant advances in NLP and is the foundation for models like BERT, GPT, and T5.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Transformer Architecture in NLP".to_string()),
                ("category".to_string(), "AI".to_string()),
                ("difficulty".to_string(), "advanced".to_string()),
                ("topics".to_string(), "transformers, NLP, attention, BERT, GPT".to_string()),
            ]),
        },
        Document {
            id: "rust_systems_programming".to_string(),
            content: "Rust is a systems programming language that runs blazingly fast, prevents segfaults, and guarantees thread safety. Rust achieves memory safety without garbage collection, and reference counting is optional. Rust was originally developed by Mozilla Research as a language for developing the Servo browser engine. The language has gained adoption in systems programming and web assembly. Rust's rich type system and ownership model guarantee memory-safety and thread-safety while maintaining zero-cost abstractions.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Rust Systems Programming".to_string()),
                ("category".to_string(), "Programming".to_string()),
                ("difficulty".to_string(), "intermediate".to_string()),
                ("topics".to_string(), "rust, systems programming, memory safety".to_string()),
            ]),
        },
        Document {
            id: "vector_databases".to_string(),
            content: "Vector databases are specialized databases designed to store, index, and search high-dimensional vector data efficiently. They are essential for applications involving machine learning, particularly for similarity search, recommendation systems, and retrieval-augmented generation (RAG). Vector databases use various indexing techniques like LSH (Locality-Sensitive Hashing), IVF (Inverted File), and HNSW (Hierarchical Navigable Small World) to enable fast approximate nearest neighbor search. Popular vector databases include Pinecone, Weaviate, Milvus, and Chroma.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Vector Databases and Similarity Search".to_string()),
                ("category".to_string(), "Database".to_string()),
                ("difficulty".to_string(), "intermediate".to_string()),
                ("topics".to_string(), "vector databases, similarity search, embeddings, indexing".to_string()),
            ]),
        },
        Document {
            id: "rag_systems".to_string(),
            content: "Retrieval-Augmented Generation (RAG) is an AI framework that combines information retrieval with text generation. RAG systems first retrieve relevant information from a knowledge base or document collection, then use that retrieved context to generate more accurate and informative responses. This approach addresses limitations of pure generative models by grounding responses in factual, up-to-date information. RAG systems typically involve document chunking, embedding generation, vector search, and context integration with language models.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Retrieval-Augmented Generation Systems".to_string()),
                ("category".to_string(), "AI".to_string()),
                ("difficulty".to_string(), "advanced".to_string()),
                ("topics".to_string(), "RAG, retrieval, generation, LLM, context".to_string()),
            ]),
        },
        Document {
            id: "distributed_computing".to_string(),
            content: "Distributed computing is a field of computer science that studies distributed systems. A distributed system is a system whose components are located on different networked computers, which communicate and coordinate their actions by passing messages to one another. The components interact with each other in order to achieve a common goal. Distributed computing enables processing large datasets across multiple machines, provides fault tolerance through redundancy, and allows for horizontal scaling. Key challenges include network partitions, consistency models, and coordination protocols.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Distributed Computing Systems".to_string()),
                ("category".to_string(), "Computer Science".to_string()),
                ("difficulty".to_string(), "advanced".to_string()),
                ("topics".to_string(), "distributed systems, networking, scalability, fault tolerance".to_string()),
            ]),
        },
        Document {
            id: "quantum_computing".to_string(),
            content: "Quantum computing is a type of computation that harnesses the collective properties of quantum states, such as superposition, interference, and entanglement, to perform calculations. The devices that perform quantum computations are known as quantum computers. Quantum computers have the potential to solve certain computational problems significantly faster than classical computers. Current applications include cryptography, optimization, machine learning, and simulation of quantum systems. However, practical quantum computers face challenges with decoherence and error rates.".to_string(),
            metadata: HashMap::from([
                ("title".to_string(), "Introduction to Quantum Computing".to_string()),
                ("category".to_string(), "Computer Science".to_string()),
                ("difficulty".to_string(), "advanced".to_string()),
                ("topics".to_string(), "quantum computing, qubits, superposition, quantum algorithms".to_string()),
            ]),
        },
    ]
}

use rtx_transformers::rag::integration_example::{GenerationConfig, RAGGenerationPipeline};