image.rs

use std::{collections::BTreeMap, io::IsTerminal};

use secrecy::{ExposeSecret, SecretString};
use serde::Serialize;
use snafu::{ResultExt, Snafu};

use crate::{
    cli::{Format, ImageCheckArguments, ImageListArguments, ImageSizeArguments, Pretty},
    config::Config,
    core::bakefile::{self, Targets, TargetsOptions},
    models::{Manifest, TagList},
    utils::{
        format_image_index_manifest_tag, format_image_manifest_tag,
        format_registry_token_env_var_name,
    },
};

#[derive(Debug, Snafu)]
pub enum Error {
    #[snafu(display("failed to serialize list as JSON"))]
    SerializeList { source: serde_json::Error },

    #[snafu(display("failed to build list of targets"))]
    BuildTargets { source: bakefile::TargetsError },

    #[snafu(display("failed to build request client"))]
    BuildClient { source: reqwest::Error },

    #[snafu(display("failed to send request ({url})"))]
    SendRequest { source: reqwest::Error, url: String },

    #[snafu(display("failed to deserialize response"))]
    DeserializeResponse { source: reqwest::Error },

    #[snafu(display("missing images found: \n {missing_images}", missing_images = missing_images.join("\n")))]
    MissingVersion { missing_images: Vec<String> },
}

/// This is the `boil show images` command handler function.
pub fn list_images(arguments: ImageListArguments) -> Result<(), Error> {
    let targets = if arguments.image.is_empty() {
        Targets::all(TargetsOptions {
            only_entry: true,
            non_recursive: false,
        })
        .context(BuildTargetsSnafu)?
    } else {
        Targets::set(
            &arguments.image,
            TargetsOptions {
                only_entry: true,
                non_recursive: false,
            },
        )
        .context(BuildTargetsSnafu)?
    };

    let list: BTreeMap<String, Vec<String>> = targets
        .into_iter()
        .map(|(image_name, (image_config, _))| {
            let versions: Vec<_> = image_config
                .versions
                .into_iter()
                .map(|(image_version, _)| image_version)
                .collect();
            (image_name, versions)
        })
        .collect();

    serialize_to_stdout(list, arguments.pretty)
}

pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Result<(), Error> {
    let targets = if arguments.image.is_empty() {
        Targets::all(TargetsOptions {
            only_entry: true,
            non_recursive: true,
        })
        .context(BuildTargetsSnafu)?
    } else {
        Targets::set(
            &arguments.image,
            TargetsOptions {
                only_entry: true,
                non_recursive: true,
            },
        )
        .context(BuildTargetsSnafu)?
    };

    let mut registry_tokens = BTreeMap::new();
    let client = reqwest::ClientBuilder::new()
        .build()
        .context(BuildClientSnafu)?;

    let mut missing_images = Vec::new();

    for (image_name, (image_config, _)) in targets {
        for (registry, registry_options) in image_config.metadata.registries {
            // Add tokens to a map so that we don't need construct the key and retrieve the value
            // over and over again.
            let registry_token = registry_tokens.entry(registry.clone()).or_insert_with(|| {
                let name = format_registry_token_env_var_name(&registry);
                std::env::var(name).ok().map(SecretString::from)
            });

            println!(
                "Checking {registry}/{registry_namespace}/{image_name}",
                registry_namespace = registry_options.namespace,
            );

            let url = format!(
                "https://{registry}/v2/{registry_namespace}/{image_name}/tags/list",
                registry_namespace = registry_options.namespace,
            );
            let request = client.get(&url);

            let request = match &registry_token {
                Some(registry_token) => request.bearer_auth(registry_token.expose_secret()),
                None => request,
            };

            let tag_list: TagList = request
                .send()
                .await
                .context(SendRequestSnafu { url })?
                .json()
                .await
                .context(DeserializeResponseSnafu)?;

            for (image_version, _) in image_config.versions.iter() {
                let index_manifest_tag = format_image_index_manifest_tag(
                    image_version,
                    &config.metadata.vendor_tag_prefix,
                    &arguments.image_version,
                );

                let image_identifier = format!("{image_name}:{index_manifest_tag}");

                println!("- {image_identifier}");
                if !tag_list.tags.contains(&index_manifest_tag) {
                    missing_images.push(image_identifier);
                }
            }
        }
    }

    if !missing_images.is_empty() {
        return MissingVersionSnafu { missing_images }.fail();
    }

    Ok(())
}

pub async fn calculate_size(arguments: ImageSizeArguments, config: Config) -> Result<(), Error> {
    let targets = if arguments.image.is_empty() {
        Targets::all(TargetsOptions {
            only_entry: true,
            non_recursive: true,
        })
        .context(BuildTargetsSnafu)?
    } else {
        Targets::set(
            &arguments.image,
            TargetsOptions {
                only_entry: true,
                non_recursive: true,
            },
        )
        .context(BuildTargetsSnafu)?
    };

    let mut registry_tokens = BTreeMap::new();
    let client = reqwest::ClientBuilder::new()
        .build()
        .context(BuildClientSnafu)?;

    #[derive(Serialize)]
    struct SizeResult {
        images: BTreeMap<String, u64>,
        total: u64,
    }

    let mut result = SizeResult {
        images: BTreeMap::new(),
        total: 0,
    };

    for (image_name, (image_config, _)) in targets {
        for (registry, registry_options) in image_config.metadata.registries {
            // Add tokens to a map so that we don't need construct the key and retrieve the value
            // over and over again.
            let registry_token = registry_tokens.entry(registry.clone()).or_insert_with(|| {
                let name = format_registry_token_env_var_name(&registry);
                std::env::var(name).ok().map(SecretString::from)
            });

            for (image_version, _) in image_config.versions.iter() {
                let image_index_manifest_tag = format_image_index_manifest_tag(
                    image_version,
                    &config.metadata.vendor_tag_prefix,
                    &arguments.vendor_version,
                );

                let manifest_tag = format_image_manifest_tag(
                    &image_index_manifest_tag,
                    arguments.target_platform.architecture(),
                    // Never strip the architecture, because we need to reference the exact manifest
                    // to be able to calculate the sizes
                    false,
                );

                let url = format!(
                    "https://{registry}/v2/{registry_namespace}/{image_name}/manifests/{manifest_tag}",
                    registry_namespace = registry_options.namespace,
                );

                let request = client.get(&url);

                let request = match &registry_token {
                    Some(registry_token) => request.bearer_auth(registry_token.expose_secret()),
                    None => request,
                };

                let manifest: Manifest = request
                    .send()
                    .await
                    .context(SendRequestSnafu { url })?
                    .json()
                    .await
                    .context(DeserializeResponseSnafu)?;

                let layer_size = manifest.layers.iter().fold(0u64, |acc, e| acc + e.size);

                let size = result.images.entry(image_name.clone()).or_default();
                *size += layer_size;

                result.total += layer_size;
            }
        }
    }

    match arguments.format {
        Format::Plain => {
            if let Some(max_width) = result
                .images
                .iter()
                .map(|i| i.0.len())
                .max_by(|lhs, rhs| lhs.cmp(rhs))
            {
                // let sorted = result.images;
                for (image, size) in result.images {
                    println!(
                        "{image:max_width$} {size}",
                        size = humansize::format_size(size, humansize::BINARY)
                    );
                }

                println!(
                    "{total_text:max_width$} {total}",
                    total_text = "Total",
                    total = humansize::format_size(result.total, humansize::BINARY)
                )
            }

            Ok(())
        }
        Format::Json => serialize_to_stdout(&result, arguments.pretty),
    }
}

fn serialize_to_stdout<T>(value: T, pretty: Pretty) -> Result<(), Error>
where
    T: Serialize,
{
    let stdout = std::io::stdout();

    match pretty {
        Pretty::Always | Pretty::Auto if stdout.is_terminal() => {
            serde_json::to_writer_pretty(stdout, &value)
        }
        _ => serde_json::to_writer(stdout, &value),
    }
    .context(SerializeListSnafu)
}