use std::time::Instant;

use rand::{thread_rng, Rng};
use tracing::{metadata::LevelFilter, subscriber::set_global_default};
use tracing_subscriber::{fmt::format::FmtSpan, layer::SubscriberExt, EnvFilter, Layer, Registry};
use vectordb::Vector;

fn init_tracing() -> Result<(), Box<dyn std::error::Error>> {
    let format_layer = tracing_subscriber::fmt::layer()
        .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
        .with_filter(
            EnvFilter::builder()
                .with_default_directive(LevelFilter::INFO.into())
                .from_env_lossy(),
        );

    let subscriber = Registry::default().with(format_layer);

    set_global_default(subscriber)?;

    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    init_tracing()?;

    let db = vectordb::VectorDb::open("./vectordb", 2)?;

    let total_chunks = 32;
    let chunk_size = 1024 * 8;

    for _ in 0..total_chunks {
        let vectors = (0..chunk_size)
            .map(|_| thread_rng().gen::<[f32; 7]>().into())
            .collect::<Vec<_>>();

        tracing::info!("Inserting vectors");
        let now = Instant::now();
        db.insert_vectors_blocking(vectors)?;
        tracing::info!("Inserting vectors - Took {:?}", now.elapsed());
    }

    let existing_vector: Vector<7> = thread_rng().gen::<[f32; 7]>().into();

    db.insert_vector_blocking(existing_vector.clone())?;

    let test_vecs = [existing_vector]
        .into_iter()
        .chain((0..5).map(|_| thread_rng().gen::<[f32; 7]>().into()))
        .collect::<Vec<_>>();

    let span = tracing::info_span!("Similarities");
    let guard = span.enter();
    for (i, v) in test_vecs.iter().enumerate() {
        let span = tracing::info_span!("test", id = %i);
        let guard = span.enter();

        let similarities = db.find_similarities_blocking(v.clone(), None, 5)?;

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            tracing::info!(
                "similar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }

        drop(guard);
        drop(span);
    }
    drop(guard);
    drop(span);

    let span = tracing::info_span!("Furthest similarities");
    let guard = span.enter();
    for (i, v) in test_vecs.iter().enumerate() {
        let span = tracing::info_span!("test", id = %i);
        let guard = span.enter();

        let similarities = db.find_furthest_similarities_blocking(v.clone(), None, 5)?;

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            tracing::info!(
                "similar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }

        drop(guard);
        drop(span);
    }
    drop(guard);
    drop(span);

    let span = tracing::info_span!("Dissimilarities");
    let guard = span.enter();
    for (i, v) in test_vecs.iter().enumerate() {
        let span = tracing::info_span!("test", id = %i);
        let guard = span.enter();

        let similarities = db.find_dissimilarities_blocking(v.clone(), None, 5)?;

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            tracing::info!(
                "dissimilar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }

        drop(guard);
        drop(span);
    }
    drop(guard);
    drop(span);

    let span = tracing::info_span!("Closest dissimilarities");
    let guard = span.enter();
    for (i, v) in test_vecs.iter().enumerate() {
        let span = tracing::info_span!("test", id = %i);
        let guard = span.enter();

        let similarities = db.find_closest_dissimilarities_blocking(v.clone(), None, 5)?;

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            tracing::info!(
                "dissimilar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }

        drop(guard);
        drop(span);
    }
    drop(guard);
    drop(span);

    Ok(())
}