use std::sync::Arc;

use futures::future::join_all;
use thingvellir::{
    service_builder, DataLoadRequest, DefaultCommitPolicy, LoadFromUpstream, ServiceData,
    ServiceHandle, ShardError, ShardStats, UpstreamError, UpstreamFactory,
};
use tokio::time::Duration;

use cassandra_cpp::{Cluster, PreparedStatement, Session};
/// This really dumb upstream just simply creates structs using
/// their default impl when requested.

struct TestCassandraUpstreamFactory {
    session: Arc<Session>,
}

struct TestCounterUpstream {
    session: Arc<Session>,
    select_statement: PreparedStatement,
}

impl LoadFromUpstream<i64, Counter> for TestCounterUpstream {
    fn load(&mut self, request: DataLoadRequest<i64, Counter>) {
        let mut statement = self.select_statement.bind();
        statement.bind_int64(0, *request.key()).unwrap();
        let fut = self.session.execute(&statement);
        request.spawn_default(async move {
            let result = fut.await?;
            let row = result.first_row().ok_or(UpstreamError::KeyNotFound)?;
            let s = row.get_column(0)?.get_i64()?;

            Ok(Counter { s })
        });
    }
}

impl UpstreamFactory<i64, Counter> for TestCassandraUpstreamFactory {
    type Upstream = TestCounterUpstream;

    fn create(&mut self) -> TestCounterUpstream {
        let session = self.session.clone();
        let select_statement = session
            .prepare("SELECT count FROM counter WHERE id = ?")
            .unwrap()
            .wait()
            .unwrap();

        TestCounterUpstream {
            session,
            select_statement,
        }
    }
}

#[derive(Default, Clone)]
struct Counter {
    s: i64,
}

impl Counter {
    fn incr(&mut self) -> i64 {
        self.s += 1;
        self.s
    }

    fn get(&self) -> i64 {
        self.s
    }
}

impl ServiceData for Counter {}

#[tokio::main]
async fn main() -> Result<(), ShardError> {
    let mut cluster = Cluster::default();
    cluster.set_contact_points("127.0.0.1").unwrap();
    let session = Session::new();
    session.connect_keyspace(&cluster, "test_keyspace")?.await?;

    let mut service =
        service_builder::<i64, Counter>(1000000).build(TestCassandraUpstreamFactory {
            session: Arc::new(session),
        });

    async fn bench(h: ServiceHandle<i64, Counter>) {
        let mut futs = vec![];
        for i in 0..50 {
            let s = i * 2000;
            let e = s + 2000;
            let mut handle = h.handle();
            let fut = tokio::spawn(async move {
                for i in s..e {
                    let result = handle.execute(i, |n| n.get()).await.unwrap();
                    // println!("Got result {:?}", result);
                }
            });

            futs.push(fut);
        }
        join_all(futs).await;
    }

    for n in 1..5 {
        let s = std::time::Instant::now();
        bench(service.handle()).await;
        let e = s.elapsed();
        println!("bench {}: {:?}", n, e);

        let stats = service.get_shard_stats().await?;
        println!("stats {}: {:?}", n, stats);
        let merged = ShardStats::merge_stats(stats);
        println!("merged {}: {:?}", n, merged);
    }

    Ok(())
}
