use bytes::{Bytes, BytesMut};
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rand::Rng;
use sia_core::rhp4::SECTOR_SIZE;
use sia_core::signing::PrivateKey;
use sia_core::types::v2::NetAddress;
use sia_storage::mock::{MockDownloader, MockHosts, MockUploader};
use sia_storage::{AppKey, DownloadOptions, Host, Object, UploadOptions};
use std::io::Cursor;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, sink};
use tokio::runtime;

async fn upload_object(uploader: Arc<MockUploader>, input: Bytes, opts: UploadOptions) -> Object {
    let r = Cursor::new(input);
    uploader
        .upload(Object::default(), r, opts)
        .await
        .expect("upload failed")
}

fn upload_benchmark(c: &mut Criterion) {
    let _ = env_logger::builder().is_test(true).try_init();
    let runtime = runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("failed to create global runtime");

    let hosts = MockHosts::new();
    hosts.update(
        (0..90)
            .map(|_| Host {
                public_key: PrivateKey::from_seed(&rand::random()).public_key(),
                addresses: vec![NetAddress {
                    protocol: sia_core::types::v2::Protocol::QUIC,
                    address: "localhost:1234".to_string(),
                }],
                country_code: "US".to_string(),
                latitude: 0.0,
                longitude: 0.0,
                good_for_upload: true,
            })
            .collect(),
        true,
    );

    let app_key = Arc::new(AppKey::import(rand::random()));
    let uploader = Arc::new(MockUploader::new(hosts.clone(), app_key.clone()));
    let downloader = Arc::new(MockDownloader::new(hosts.clone(), app_key.clone()));
    let mut input = BytesMut::zeroed(SECTOR_SIZE * 30); // 3 full slabs
    rand::rng().fill_bytes(&mut input);
    let input = input.freeze();

    let mut large_group = c.benchmark_group("120MiB");
    large_group.throughput(Throughput::Bytes(input.len() as u64));

    // all shards in flight
    large_group.bench_with_input(
        BenchmarkId::new("upload", "90 inflight"),
        &input,
        |b, input| {
            b.to_async(&runtime).iter(|| async {
                upload_object(
                    uploader.clone(),
                    input.clone(),
                    UploadOptions {
                        max_inflight: 90,
                        ..Default::default()
                    },
                )
                .await;
            });
            hosts.clear();
        },
    );

    large_group.bench_with_input(
        BenchmarkId::new("upload", "10 inflight"),
        &input,
        |b, input| {
            b.to_async(&runtime).iter(|| async {
                upload_object(
                    uploader.clone(),
                    input.clone(),
                    UploadOptions {
                        max_inflight: 10,
                        ..Default::default()
                    },
                )
                .await;
            });
            hosts.clear();
        },
    );

    large_group.bench_with_input(BenchmarkId::new("upload", "default"), &input, |b, input| {
        b.to_async(&runtime).iter(|| async {
            upload_object(uploader.clone(), input.clone(), UploadOptions::default()).await;
        });
        hosts.clear();
    });

    let object = runtime.block_on(async {
        upload_object(uploader.clone(), input.clone(), UploadOptions::default()).await
    });

    large_group.bench_with_input(
        BenchmarkId::new("download", "30 inflight"),
        &object,
        |b, object| {
            b.to_async(&runtime).iter(|| async {
                let mut reader = downloader
                    .download(
                        object,
                        DownloadOptions {
                            max_inflight: 30,
                            ..Default::default()
                        },
                    )
                    .unwrap();
                tokio::io::copy(&mut reader, &mut sink())
                    .await
                    .expect("download to complete");
            });
        },
    );

    large_group.bench_with_input(
        BenchmarkId::new("download", "10 inflight"),
        &object,
        |b, object| {
            b.to_async(&runtime).iter(|| async {
                let mut reader = downloader
                    .download(
                        object,
                        DownloadOptions {
                            max_inflight: 10,
                            ..Default::default()
                        },
                    )
                    .unwrap();
                tokio::io::copy(&mut reader, &mut sink())
                    .await
                    .expect("download to complete");
            });
        },
    );

    large_group.bench_with_input(
        BenchmarkId::new("download", "default"),
        &object,
        |b, object| {
            b.to_async(&runtime).iter(|| async {
                let mut reader = downloader
                    .download(object, DownloadOptions::default())
                    .unwrap();
                tokio::io::copy(&mut reader, &mut sink())
                    .await
                    .expect("download to complete");
            });
        },
    );

    large_group.finish();

    let mut ttfb_group = c.benchmark_group("ttfb");

    ttfb_group.bench_function("120MiB", |b| {
        b.to_async(&runtime).iter_custom(|iters| {
            let downloader = downloader.clone();
            let object = object.clone();
            async move {
                let mut total = std::time::Duration::ZERO;
                for _ in 0..iters {
                    let start = std::time::Instant::now();
                    let mut reader = downloader
                        .download(&object, DownloadOptions::default())
                        .unwrap();
                    let mut buf = [0u8; 1];
                    reader.read(&mut buf).await.expect("read to succeed");
                    total += start.elapsed();
                }
                total
            }
        });
    });

    ttfb_group.bench_function("64B", |b| {
        b.to_async(&runtime).iter_custom(|iters| {
            let downloader = downloader.clone();
            let object = object.clone();
            async move {
                let mut total = std::time::Duration::ZERO;
                for _ in 0..iters {
                    let start = std::time::Instant::now();
                    let mut reader = downloader
                        .download(
                            &object,
                            DownloadOptions {
                                length: Some(64),
                                ..Default::default()
                            },
                        )
                        .unwrap();
                    let mut buf = [0u8; 1];
                    reader.read(&mut buf).await.expect("read to succeed");
                    total += start.elapsed();
                }
                total
            }
        });
    });
    ttfb_group.finish();
}

criterion_group!(benches, upload_benchmark);
criterion_main!(benches);
