use stopwatch::Stopwatch;
use tokio::io::AsyncWriteExt;
use tokio::process;

#[tokio::main]
async fn main() -> eyre::Result<()> {
    tracing_subscriber::fmt().pretty().init();

    let client = dagger_sdk::connect().await?;

    let (tx, mut rx) = tokio::sync::mpsc::channel(4);

    async_scoped::TokioScope::scope_and_block(|s| {
        let dclient = client.clone();
        let dtx = tx.clone();
        s.spawn(async move {
            let dclient = dclient.pipeline("build-image");
            tracing::info!("building 1. complete image");

            let sw = Stopwatch::start_new();
            let src = dclient
                .host()
                .directory_opts(
                    "bench_app",
                    dagger_sdk::HostDirectoryOptsBuilder::default()
                        .exclude(vec!["target/"])
                        .build()
                        .unwrap(),
                )
                .id()
                .await
                .unwrap();

            let _ = dclient
                .container()
                .from("rustlang/rust:nightly")
                .with_directory("/mnt/app", src)
                .with_workdir("/mnt/app")
                .with_exec(vec!["cargo", "build", "--release"])
                .with_exec(vec!["target/release/bench_app", "bench"])
                .exit_code()
                .await
                .unwrap();

            dtx.send(format!(
                "1. build-image: {}",
                sw.elapsed_ms() as f64 / 1000.0
            ))
            .await
            .unwrap();
        });

        let dclient = client.clone();
        let dtx = tx.clone();
        s.spawn(async move {
            let dclient = dclient.pipeline("pull-image");
            tracing::info!("building 2. pull image");
            let sw = Stopwatch::start_new();

            if let Err(e) = process::Command::new("docker")
                .args(&[
                    "image",
                    "rm",
                    "kasperhermansen/dagger-runtime-benchmark:latest",
                ])
                .output()
                .await
            {
                tracing::warn!(msg = e.to_string(), "failed to delete image");
            }

            let _ = dclient
                .container()
                .from("kasperhermansen/dagger-runtime-benchmark:latest")
                .with_exec(vec!["/usr/bin/bench_app", "bench"])
                .exit_code()
                .await
                .unwrap();
            dtx.send(format!(
                "2. pull-image: {}",
                sw.elapsed_ms() as f64 / 1000.0
            ))
            .await
            .unwrap();
        });

        let dclient = client.clone();
        let dtx = tx.clone();
        s.spawn(async move {
            let dclient = dclient.pipeline("run-wasm");
            tracing::info!("building 3. pull wasm");
            let sw = Stopwatch::start_new();

            let src = dclient
                .host()
                .directory(".")
                .file("dist/wasm/bench_app.wasm")
                .id()
                .await
                .unwrap();

            if let Err(e) = process::Command::new("docker")
                .args(&["image", "rm", "kasperhermansen/dagger-wasmer:latest"])
                .output()
                .await
            {
                tracing::warn!(msg = e.to_string(), "failed to delete image");
            }

            let _ = dclient
                .container()
                .from("kasperhermansen/dagger-wasmer:latest")
                .with_file("/usr/bin/bench_app.wasm", src)
                .with_exec(vec!["wasmer", "run", "/usr/bin/bench_app.wasm", "bench"])
                .exit_code()
                .await
                .unwrap();

            dtx.send(format!("3. run-wasm: {}", sw.elapsed_ms() as f64 / 1000.0))
                .await
                .unwrap();
        });

        let dclient = client.clone();
        let dtx = tx.clone();
        s.spawn(async move {
            let dclient = dclient.pipeline("run-binary");
            tracing::info!("building 4. run binary");
            let sw = Stopwatch::start_new();

            let src = dclient
                .host()
                .directory(".")
                .file("dist/binary/bench_app")
                .id()
                .await
                .unwrap();

            if let Err(e) = process::Command::new("docker")
                .args(&["image", "rm", "debian:bullseye-slim"])
                .output()
                .await
            {
                tracing::warn!(msg = e.to_string(), "failed to delete image");
            }

            let _ = dclient
                .container()
                .from("debian:bullseye-slim")
                .with_file("/usr/bin/bench_app", src)
                .with_exec(vec!["/usr/bin/bench_app", "bench"])
                .exit_code()
                .await
                .unwrap();

            dtx.send(format!(
                "4. run-binary: {}",
                sw.elapsed_ms() as f64 / 1000.0
            ))
            .await
            .unwrap();
        });
    });

    drop(tx);

    let mut file = tokio::fs::File::create("reports/timings.md").await?;
    file.write_all(b"# Dagger runtime benchmark results\n")
        .await?;

    let mut results = Vec::new();
    while let Some(res) = rx.recv().await {
        results.push(res);
    }
    results.sort();

    for res in results {
        file.write_all(format!("{}\n", res).as_bytes()).await?;
    }

    file.sync_all().await?;

    Ok(())
}