xz-rag 0.1.1

use std::sync::Arc;

use async_trait::async_trait;
use xz_rag::channels::semantic::{Embedder, SemanticSearch};
use xz_rag::pipeline::channel::{ChannelConfig, ChannelPipeline};
use xz_rag::types::chunk::ChunkMetadata;
use xz_rag::types::config::RagEngineInfo;
use xz_rag::types::rag::RagRequest;
use xz_rag::types::retrieval::RetrieveRequest;
use xz_rag::{DefaultRagEngine, RagEngine, RagError};

// ---------------------------------------------------------------------------
// Mock embedder — returns a fixed zero vector for any input
// ---------------------------------------------------------------------------

struct MockEmbedder;

#[async_trait]
impl Embedder for MockEmbedder {
    async fn embed(&self, text: &[String]) -> Result<Vec<Vec<f32>>, RagError> {
        // Return a zero vector for each input text (dimension = 4)
        Ok(text.iter().map(|_| vec![0.0_f32; 4]).collect())
    }

    fn dimensions(&self) -> usize {
        4
    }
}

// ---------------------------------------------------------------------------
// Mock semantic store — returns hard-coded chunks regardless of query
// ---------------------------------------------------------------------------

struct MockSemanticStore {
    chunks: Vec<(String, f32, ChunkMetadata, String, String)>,
}

#[async_trait]
impl SemanticSearch for MockSemanticStore {
    async fn search(
        &self,
        _query_embedding: &[f32],
        top_k: usize,
        _namespace: Option<&str>,
    ) -> Result<Vec<(String, f32, ChunkMetadata, String, String)>, RagError> {
        let mut results: Vec<_> = self.chunks.clone();
        results.truncate(top_k);
        Ok(results)
    }
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // -- Arrange mock data ------------------------------------------------

    let chunks = vec![
        (
            "chunk-1".into(),
            0.92_f32,
            ChunkMetadata {
                document_title: Some("Rust Book".into()),
                source: Some("doc/rust-book.md".into()),
                ..ChunkMetadata::default()
            },
            "Rust is a systems programming language focused on safety and performance."
                .into(),
            "doc-1".into(),
        ),
        (
            "chunk-2".into(),
            0.85_f32,
            ChunkMetadata {
                document_title: Some("Async Rust".into()),
                source: Some("doc/async-rust.md".into()),
                ..ChunkMetadata::default()
            },
            "Tokio is the primary async runtime for Rust, providing I/O, timers, and synchronization."
                .into(),
            "doc-2".into(),
        ),
        (
            "chunk-3".into(),
            0.73_f32,
            ChunkMetadata {
                document_title: Some("Rust Book".into()),
                source: Some("doc/rust-book.md".into()),
                ..ChunkMetadata::default()
            },
            "Ownership is Rust's most unique feature, enabling memory safety without a garbage collector."
                .into(),
            "doc-1".into(),
        ),
    ];

    let embedder: Arc<dyn Embedder> = Arc::new(MockEmbedder);
    let store: Arc<dyn SemanticSearch> = Arc::new(MockSemanticStore { chunks });

    // -- Build engine -----------------------------------------------------

    let pipeline = ChannelPipeline::new(vec![ChannelConfig::semantic(0.5, 10).with_min_score(0.1)]);

    let engine = DefaultRagEngine::builder()
        .name("basic-rag")
        .version("0.1.0")
        .pipeline(pipeline)
        .embedder(embedder)
        .semantic_store(store)
        .build();

    // -- Inspect engine metadata ------------------------------------------

    let info: RagEngineInfo = engine.engine_info();
    println!("=== Engine Info ===");
    println!("  Name:              {}", info.name);
    println!("  Version:           {}", info.version);
    println!("  Channels:          {:?}", info.supported_channels);
    println!("  Max context:       {} tokens", info.max_context_window);
    println!("  Streaming:         {}", info.supports_streaming);
    println!("  Reranking:         {}", info.reranking_enabled);
    println!();

    // -- 1. Retrieve only -------------------------------------------------

    let retrieve_req = RetrieveRequest::builder("What is Rust?").top_k(5).build();

    let retrieve_result = engine.retrieve(&retrieve_req).await?;

    println!("=== Retrieve Results ===");
    println!("  Query:             {}", retrieve_result.effective_query);
    println!("  Latency:           {} ms", retrieve_result.latency_ms);
    println!("  Hits:              {}", retrieve_result.hits.len());
    for hit in &retrieve_result.hits {
        println!("    [{:.4}] (ch={}) {}", hit.score, hit.channel, hit.content);
    }
    println!();

    // -- 2. Retrieve + generate (placeholder, no LLM feature) ------------

    let rag_req: RagRequest = RagRequest::builder("What makes Rust safe?")
        .system_prompt("You are a Rust expert.")
        .retrieve_config(RetrieveRequest::builder("What makes Rust safe?").top_k(3).build())
        .build();

    let rag_response = engine.retrieve_and_generate(&rag_req).await?;

    println!("=== Retrieve + Generate ===");
    println!("  Answer:");
    for line in rag_response.answer.lines() {
        println!("    {line}");
    }
    println!();
    println!("  Citations:         {}", rag_response.citations.len());
    for cit in &rag_response.citations {
        println!(
            "    [{}] {} (score={:.4})",
            cit.index,
            cit.document_title.as_deref().unwrap_or("?"),
            cit.score
        );
    }
    println!();
    println!("  Token usage:");
    println!("    Context:  {}", rag_response.usage.context_tokens);
    println!("    Prompt:   {}", rag_response.usage.prompt_tokens);
    println!("    Complete: {}", rag_response.usage.completion_tokens);
    println!("    Chunks used:   {}", rag_response.usage.chunks_used);
    println!("    Chunks dropped: {}", rag_response.usage.chunks_dropped);
    println!();
    println!("  Total latency:     {} ms", rag_response.total_latency_ms);

    Ok(())
}