scarfbench-cli 0.1.2

use crate::utils::{
    get_or_create_and_get_scarfbench_home_dir,
    progress_bar::{ProgressBar, ProgressReader},
};
use anyhow::{Context, Result};
use bon::Builder;
use clap::Args;
use flate2::bufread::GzDecoder;
use kdam::term;
use reqwest::blocking::{Client, RequestBuilder};
use serde::{Deserialize, Serialize};
use std::{
    fs::{self},
    io::{BufReader, IsTerminal},
    path::PathBuf,
};
use tar::Archive;
#[derive(Args, Debug)]
pub struct BenchPullArgs {
    #[arg(
        short,
        long,
        help = "Path to where the benchmark is to be saved. Defaults to ~/.scarfbench/benchmark",
        value_name = "DIR"
    )]
    pub benchmark_dest: Option<PathBuf>,

    #[arg(
        long,
        help = "Version of scarfbench to pull.",
        value_name = "VERSION"
    )]
    pub version: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
struct Release {
    assets: Vec<Asset>,
    tag_name: String,
}

#[derive(Debug, Serialize, Deserialize)]
struct Asset {
    name: String,
    browser_download_url: String,
}

#[derive(Debug, Builder)]
pub struct PullScarfBench {
    /// Name (prefix) of the benchmark
    #[builder(default= "benchmark-v".to_string())]
    pub asset_name_prefix: String,

    /// Version of scarfbench to pull. Default is latest.
    pub version: Option<String>,

    /// Where to download the benchmark to.
    pub dest_dir: PathBuf,
}

impl PullScarfBench {
    // These are for my CI ghactions builds...
    fn github_token() -> Option<String> {
        std::env::var("SCARF_BENCH_GITHUB_TOKEN")
            .ok()
            .filter(|v| !v.trim().is_empty())
            .or_else(|| {
                std::env::var("GITHUB_TOKEN")
                    .ok()
                    .filter(|v| !v.trim().is_empty())
            })
    }

    fn authorize_me_maybe(
        request: RequestBuilder,
        token: Option<&str>,
    ) -> RequestBuilder {
        match token {
            Some(token) => request.bearer_auth(token),
            None => request,
        }
    }

    pub fn exec(&self) -> anyhow::Result<i32> {
        let token = Self::github_token();
        let client = Client::builder().user_agent("scarf-cli").build()?;

        // Get the download URL
        let api_url = match self.version.as_deref() {
            Some(v) => format!(
                "https://api.github.com/repos/scarfbench/benchmark/releases/tags/{}",
                v
            ),
            None => format!(
                "https://api.github.com/repos/scarfbench/benchmark/releases/latest"
            ),
        };
        log::info!("Downloading from {api_url}");

        // Get releases
        let release_response = Self::authorize_me_maybe(client.get(&api_url), token.as_deref())
            .header("User-Agent", "scarf")
            .send()
            .with_context(|| format!("Unable to fetch the release metadata from {api_url}"))?
            .error_for_status()
            .with_context(|| {
                format!(
                    "GitHub API returned an error status while reading {api_url}. If this repo is private, set SCARF_BENCH_GITHUB_TOKEN (or GITHUB_TOKEN) with contents:read access to scarfbench/benchmark."
                )
            })?;

        let releases: Release =
            release_response.json().context("Failed to parse release JSON")?;

        // Get the asset to download
        let asset = releases
            .assets
            .into_iter()
            .find(|predicate| predicate.name.contains(&self.asset_name_prefix))
            .with_context(|| {
                return format!(
                    "There are no release assets that start with {} available at {}",
                    &self.asset_name_prefix, api_url
                );
            })?;

        // Create the save destination
        fs::create_dir_all(&self.dest_dir).with_context(|| {
            format!(
                "Failed to create a directory at {}",
                &self.dest_dir.to_string_lossy()
            )
        })?;

        let response = Self::authorize_me_maybe(client.get(&asset.browser_download_url), token.as_deref())
            .send()
            .with_context(|| {
                format!(
                    "Failed to download the asset from {}",
                    asset.browser_download_url
                )
            })?
            .error_for_status()
            .with_context(|| {
                format!(
                    "Asset download returned an error status from {}. If this release is private, set SCARF_BENCH_GITHUB_TOKEN (or GITHUB_TOKEN) with access to scarfbench/benchmark.",
                    asset.browser_download_url
                )
            })?;

        // Get the total size of the payload (our benchmark tar.gz we are downloading)
        let total_size =
            response.content_length().map(|s| s as usize).unwrap_or(0);

        // Set up terminal to tell kdam if we are in a active terminal (for colors and ansi stuff)
        term::init(std::io::stderr().is_terminal());

        // initialize our progress bar
        let pb = total_size.progress("Downloading scarfbench", "B");

        // Here we are encapsulating the response stream by overlaying it with the progressbar
        let pr =
            ProgressReader::new(BufReader::new(response), pb, Some(total_size));

        // Extract by streaming the response to tar directly
        let tar = GzDecoder::new(pr);
        let mut archive = Archive::new(tar);
        archive.unpack(&self.dest_dir).with_context(|| {
            format!("Failed to extract into {}", self.dest_dir.display())
        })?;
        log::info!(
            "Benchmark successfully downloaded to {}",
            &self.dest_dir.join("benchmark").display()
        );
        Ok(0)
    }
}

/// A simple list subcommand that lists all the benchmark applications as a table.
pub fn run(bench_pull_args: BenchPullArgs) -> Result<i32> {
    PullScarfBench::builder()
        .maybe_version(bench_pull_args.version)
        .dest_dir(bench_pull_args.benchmark_dest.unwrap_or_else(|| {
            get_or_create_and_get_scarfbench_home_dir().unwrap()
        }))
        .build()
        .exec()
}