harmony/harmony_composer/src/main.rs

use bollard::models::ContainerCreateBody;
use bollard::query_parameters::{
    CreateContainerOptionsBuilder, ListContainersOptionsBuilder, LogsOptions,
    RemoveContainerOptions, StartContainerOptions, WaitContainerOptions,
};
use bollard::secret::HostConfig;
use cargo_metadata::{Artifact, Message, MetadataCommand};
use clap::{Args, Parser, Subcommand};
use futures_util::StreamExt;
use instrumentation::HarmonyComposerEvent;
use log::{debug, info, log_enabled};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use tokio::fs;

mod harmony_composer_logger;
mod instrumentation;

#[derive(Parser)]
#[command(version, about, long_about = None, flatten_help = true, propagate_version = true)]
struct GlobalArgs {
    #[arg(long, default_value = ".")]
    harmony_path: String,

    #[arg(long)]
    compile_method: Option<CompileMethod>,

    #[arg(long)]
    compile_platform: Option<String>,

    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand, Clone, Debug)]
enum Commands {
    Check(CheckArgs),
    Compile,
    Deploy(DeployArgs),
    All(AllArgs),
}

#[derive(Args, Clone, Debug)]
struct CheckArgs {
    #[arg(long, default_value = "check.sh")]
    check_script_path: String,
}

#[derive(Args, Clone, Debug)]
struct DeployArgs {
    #[arg(long = "target", short = 't', default_value = "local")]
    harmony_target: HarmonyTarget,

    #[arg(long = "profile", short = 'p', default_value = "dev")]
    harmony_profile: HarmonyProfile,

    #[arg(long = "dry-run", short = 'd', default_value = "false")]
    dry_run: bool,
}

#[derive(Args, Clone, Debug)]
struct AllArgs {
    #[command(flatten)]
    check: CheckArgs,

    #[command(flatten)]
    deploy: DeployArgs,
}

#[derive(Clone, Debug, clap::ValueEnum)]
enum HarmonyTarget {
    Local,
    Remote,
}

impl std::fmt::Display for HarmonyTarget {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HarmonyTarget::Local => f.write_str("local"),
            HarmonyTarget::Remote => f.write_str("remote"),
        }
    }
}

#[derive(Clone, Debug, clap::ValueEnum)]
enum HarmonyProfile {
    Dev,
    Staging,
    Production,
}

impl std::fmt::Display for HarmonyProfile {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HarmonyProfile::Dev => f.write_str("dev"),
            HarmonyProfile::Staging => f.write_str("staging"),
            HarmonyProfile::Production => f.write_str("production"),
        }
    }
}

#[tokio::main]
async fn main() {
    harmony_composer_logger::init();
    let cli_args = GlobalArgs::parse();

    let harmony_path = Path::new(&cli_args.harmony_path)
        .try_exists()
        .expect("couldn't check if path exists");

    instrumentation::instrument(HarmonyComposerEvent::ProjectInitializationStarted).unwrap();

    let harmony_bin_path: PathBuf = match harmony_path {
        true => compile_harmony(
            cli_args.compile_method,
            cli_args.compile_platform,
            cli_args.harmony_path.clone(),
        )
        .await
        .expect("couldn't compile harmony"),
        false => todo!("implement autodetect code"),
    };

    instrumentation::instrument(HarmonyComposerEvent::ProjectInitialized).unwrap();

    match cli_args.command {
        Some(command) => match command {
            Commands::Check(args) => {
                let check_script_str =
                    format!("{}/{}", cli_args.harmony_path, args.check_script_path);

                let check_script = Path::new(&check_script_str);

                match check_script
                    .try_exists()
                    .expect("couldn't check if path exists")
                {
                    true => (),
                    false => todo!(
                        "path {check_script_str} not found. Other paths currently unsupported."
                    ),
                };

                let check_output = Command::new(check_script)
                    .spawn()
                    .expect("failed to run check script")
                    .wait_with_output()
                    .unwrap();
                info!(
                    "check stdout: {}, check stderr: {}",
                    String::from_utf8(check_output.stdout).expect("couldn't parse from utf8"),
                    String::from_utf8(check_output.stderr).expect("couldn't parse from utf8")
                );
            }
            Commands::Deploy(args) => {
                instrumentation::instrument(HarmonyComposerEvent::DeploymentStarted {
                    target: args.harmony_target.clone(),
                    profile: args.harmony_profile.clone(),
                })
                .unwrap();

                if matches!(args.harmony_profile, HarmonyProfile::Dev)
                    && !matches!(args.harmony_target, HarmonyTarget::Local)
                {
                    instrumentation::instrument(HarmonyComposerEvent::DeploymentFailed {
                        details: format!(
                            "Cannot run profile '{}' on target '{}'. Profile '{}' can run locally only.",
                            args.harmony_profile, args.harmony_target, args.harmony_profile
                        ),
                    }).unwrap();
                    return;
                }

                let use_local_k3d = match args.harmony_target {
                    HarmonyTarget::Local => true,
                    HarmonyTarget::Remote => false,
                };

                let mut command = Command::new(harmony_bin_path);
                command
                    .env("HARMONY_USE_LOCAL_K3D", format!("{use_local_k3d}"))
                    .env("HARMONY_PROFILE", format!("{}", args.harmony_profile))
                    .env("HARMONY_DRY_RUN", format!("{}", args.dry_run))
                    .arg("-y")
                    .arg("-a");

                info!("{:?}", command);

                let deploy = command.spawn().expect("failed to run harmony deploy");

                let deploy_output = deploy.wait_with_output().unwrap();
                debug!("{}", String::from_utf8(deploy_output.stdout).unwrap());

                instrumentation::instrument(HarmonyComposerEvent::DeploymentCompleted).unwrap();
            }
            Commands::All(_args) => todo!(
                "take all previous match arms and turn them into separate functions, and call them all one after the other"
            ),
            Commands::Compile => return,
        },
        None => todo!("run interactively, ask for info on CLI"),
    }

    instrumentation::instrument(HarmonyComposerEvent::Shutdown).unwrap();
}

#[derive(Clone, Debug, clap::ValueEnum)]
enum CompileMethod {
    LocalCargo,
    Docker,
}

async fn compile_harmony(
    method: Option<CompileMethod>,
    platform: Option<String>,
    harmony_location: String,
) -> Result<PathBuf, String> {
    let platform = match platform {
        Some(p) => p,
        None => current_platform::CURRENT_PLATFORM.to_string(),
    };

    let cargo_exists = Command::new("which")
        .arg("cargo")
        .stdout(Stdio::null())
        .status()
        .expect("couldn't get `which cargo` status")
        .success();

    let method = match method {
        Some(m) => m,
        None => {
            if cargo_exists {
                CompileMethod::LocalCargo
            } else {
                CompileMethod::Docker
            }
        }
    };

    let path = match method {
        CompileMethod::LocalCargo => {
            instrumentation::instrument(HarmonyComposerEvent::ProjectCompilationStarted {
                details: "compiling project with cargo".to_string(),
            })
            .unwrap();

            compile_cargo(platform, harmony_location).await
        }
        CompileMethod::Docker => {
            instrumentation::instrument(HarmonyComposerEvent::ProjectCompilationStarted {
                details: "compiling project with docker".to_string(),
            })
            .unwrap();

            compile_docker(platform, harmony_location).await
        }
    };

    match path {
        Ok(path) => {
            instrumentation::instrument(HarmonyComposerEvent::ProjectCompiled).unwrap();
            Ok(path)
        }
        Err(err) => {
            instrumentation::instrument(HarmonyComposerEvent::ProjectCompilationFailed {
                details: err.clone(),
            })
            .unwrap();
            Err(err)
        }
    }
}

// TODO: make sure this works with cargo workspaces
async fn compile_cargo(platform: String, harmony_location: String) -> Result<PathBuf, String> {
    let metadata = MetadataCommand::new()
        .manifest_path(format!("{}/Cargo.toml", harmony_location))
        .exec()
        .unwrap();

    let stderr = if log_enabled!(log::Level::Debug) {
        Stdio::inherit()
    } else {
        Stdio::piped()
    };

    let mut cargo_build = Command::new("cargo")
        .current_dir(&harmony_location)
        .args(vec![
            "build",
            format!("--target={}", platform).as_str(),
            "--release",
            "--message-format=json-render-diagnostics",
        ])
        .stdout(Stdio::piped())
        .stderr(stderr)
        .spawn()
        .expect("run cargo command failed");

    let mut artifacts: Vec<Artifact> = vec![];
    let reader = std::io::BufReader::new(cargo_build.stdout.take().unwrap());
    for message in cargo_metadata::Message::parse_stream(reader) {
        match message.unwrap() {
            Message::CompilerMessage(_msg) => (),
            Message::CompilerArtifact(artifact) => {
                if artifact.manifest_path
                    == metadata
                        .root_package()
                        .expect("failed to get root package")
                        .manifest_path
                {
                    debug!("{:?}", artifact);
                    artifacts.push(artifact);
                }
            }
            Message::BuildScriptExecuted(_script) => (),
            Message::BuildFinished(finished) => {
                debug!("{:?}", finished);
            }
            _ => (), // Unknown message
        }
    }

    let res = cargo_build.wait(); //.expect("run cargo command failed");
    if res.is_err() {
        return Err("cargo build failed".into());
    }

    let bin = artifacts
        .last()
        .expect("no binaries built")
        .filenames
        .first()
        .expect("couldn't get filename");

    let bin_out;
    if let Some(ext) = bin.extension() {
        bin_out = PathBuf::from(format!("{}/harmony.{}", harmony_location, ext));
        let _copy_res = fs::copy(&bin, &bin_out).await;
    } else {
        bin_out = PathBuf::from(format!("{}/harmony", harmony_location));
        let _copy_res = fs::copy(&bin, &bin_out).await;
    }

    Ok(bin_out)
}

async fn compile_docker(platform: String, harmony_location: String) -> Result<PathBuf, String> {
    let docker_client =
        bollard::Docker::connect_with_local_defaults().expect("couldn't connect to docker");

    let mut filters = HashMap::new();
    filters.insert(String::from("name"), vec![String::from("harmony_build")]);
    let list_containers_options = ListContainersOptionsBuilder::new()
        .all(true)
        .filters(&filters)
        .build();

    let containers = &docker_client
        .list_containers(Some(list_containers_options))
        .await
        .expect("list containers failed");

    if !containers.is_empty() {
        docker_client
            .remove_container("harmony_build", None::<RemoveContainerOptions>)
            .await
            .expect("failed to remove container");
    }

    let options = Some(
        CreateContainerOptionsBuilder::new()
            .name("harmony_build")
            .build(),
    );

    let config = ContainerCreateBody {
        image: Some("hub.nationtech.io/harmony/harmony_composer".to_string()),
        working_dir: Some("/mnt".to_string()),
        cmd: Some(vec![
            format!("--compile-platform={}", platform),
            format!("--harmony-path=/mnt"),
            "compile".to_string(),
        ]),
        host_config: Some(HostConfig {
            binds: Some(vec![format!("{}:/mnt", harmony_location)]),
            ..Default::default()
        }),
        ..Default::default()
    };

    docker_client
        .create_container(options, config)
        .await
        .expect("couldn't create container");

    docker_client
        .start_container("harmony_build", None::<StartContainerOptions>)
        .await
        .expect("couldn't start container");

    let wait_options = WaitContainerOptions {
        condition: "not-running".to_string(),
    };
    let mut wait = docker_client
        .wait_container("harmony_build", Some(wait_options))
        .boxed();

    let mut logs_stream = docker_client.logs(
        "harmony_build",
        Some(LogsOptions {
            follow: true,
            stdout: true,
            stderr: true,
            tail: "all".to_string(),
            ..Default::default()
        }),
    );

    while let Some(l) = logs_stream.next().await {
        let l_str = l.expect("couldn't unwrap logoutput").to_string();
        println!("{}", l_str);
    }

    // wait until container is no longer running
    while (wait.next().await).is_some() {}

    // hack that should be cleaned up
    if platform.contains("windows") {
        Ok(PathBuf::from(format!("{}/harmony.exe", harmony_location)))
    } else {
        Ok(PathBuf::from(format!("{}/harmony", harmony_location)))
    }
}