use std::time::Duration;

use stratum_embeddings::{
    EmbeddingOptions, EmbeddingService, FastEmbedProvider, MemoryCache, SurrealDbStore,
    VectorStore, VectorStoreConfig,
};
use tracing::info;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt::init();

    info!("Initializing embedding service...");
    let provider = FastEmbedProvider::small()?;
    let cache = MemoryCache::new(1000, Duration::from_secs(300));
    let service = EmbeddingService::new(provider).with_cache(cache);

    info!("Creating SurrealDB in-memory store...");
    let config = VectorStoreConfig::new(384);
    let store = SurrealDbStore::new_memory("concepts", config).await?;

    let concepts = vec![
        ("ownership", "Rust's ownership system prevents memory leaks"),
        (
            "borrowing",
            "Borrowing allows references without ownership transfer",
        ),
        ("lifetimes", "Lifetimes ensure references remain valid"),
        ("traits", "Traits define shared behavior across types"),
        ("generics", "Generics enable code reuse with type safety"),
    ];

    info!("Embedding and storing {} concepts...", concepts.len());
    let options = EmbeddingOptions::default_with_cache();

    for (id, description) in &concepts {
        let embedding = service.embed(description, &options).await?;
        let metadata = serde_json::json!({
            "concept": id,
            "description": description,
            "language": "rust"
        });
        store.upsert(id, &embedding, metadata).await?;
    }

    info!("Concepts stored successfully");

    info!("Performing knowledge graph search...");
    let query = "How does Rust manage memory?";
    let query_embedding = service.embed(query, &options).await?;

    let results = store.search(&query_embedding, 3, None).await?;

    info!("Most relevant concepts for: '{}'", query);
    for (i, result) in results.iter().enumerate() {
        let concept = result.metadata["concept"].as_str().unwrap_or("N/A");
        let description = result.metadata["description"].as_str().unwrap_or("N/A");
        info!("  {}. {} [score: {:.4}]", i + 1, concept, result.score);
        info!("     {}", description);
    }

    let count = store.count().await?;
    info!("Total concepts in graph: {}", count);

    Ok(())
}