pre_game.rs

use std::path::Path;

use clap::Args;
use color_eyre::{
    Result,
    eyre::{WrapErr, bail},
};

use argument_parsers::RobotAddress;
use indicatif::ProgressBar;
use repository::{Repository, upload::get_binary};
use robot::{Network, Robot, SystemctlAction};
use tempfile::tempdir;

use crate::{
    cargo::{self, CargoCommand, build, cargo, environment::EnvironmentArguments},
    deploy_config::DeployConfig,
    progress_indicator::ProgressIndicator,
};

#[derive(Args)]
#[group(skip)]
pub struct Arguments {
    #[command(flatten)]
    pub pre_game: PreGameArguments,
    #[command(flatten)]
    pub environment: EnvironmentArguments,
    #[command(flatten, next_help_heading = "Cargo Options")]
    pub build: build::Arguments,
}

#[derive(Args)]
pub struct PreGameArguments {
    /// Do not build before uploading
    #[arg(long, visible_alias = "nich-bauen")]
    pub no_build: bool,
    /// Do not restart HULK service after uploading
    #[arg(long, visible_alias = "nich-neustarten")]
    pub no_restart: bool,
    /// Do not remove existing remote files during uploading
    #[arg(long, visible_alias = "nich-säubern")]
    pub no_clean: bool,
    /// Skip the OS version check
    #[arg(long, visible_alias = "überspringe-bs-prüfung")]
    pub skip_os_check: bool,
    /// Prepare everything for the upload without performing the actual one
    #[arg(long, visible_alias = "vorbereit")]
    pub prepare: bool,
    /// The Robots to apply the pregame to, queried from the deploy.toml if not specified
    pub robots: Option<Vec<RobotAddress>>,
}

pub async fn pre_game(arguments: Arguments, repository: &Repository) -> Result<()> {
    let config = DeployConfig::read_from_file(repository)
        .await
        .wrap_err("failed to read deploy config from file")?;

    let playing_robots = config.playing_robots()?;
    let robots = if let Some(robots) = &arguments.pre_game.robots {
        for robot in robots {
            if !playing_robots.contains(robot) {
                bail!("Robot with IP {robot} is not one of the playing Robots in the deploy.toml");
            }
        }
        robots
    } else {
        &playing_robots
    };
    let wifi = config.wifi;

    config
        .configure_repository(repository)
        .await
        .wrap_err("failed to configure repository")?;

    let upload_directory = tempdir().wrap_err("failed to get temporary directory")?;
    const BINARY_NAME: &str = "hulk_ros_z";
    let hulk_binary = get_binary(arguments.build.profile(), BINARY_NAME);

    let cargo_arguments = cargo::Arguments {
        manifest: Some(
            repository
                .root
                .join(format!("crates/{BINARY_NAME}/Cargo.toml"))
                .into_os_string(),
        ),
        environment: arguments.environment,
        cargo: arguments.build,
    };

    if !arguments.pre_game.no_build {
        cargo(cargo_arguments, repository, &[&hulk_binary])
            .await
            .wrap_err("failed to build")?;
    }
    if arguments.pre_game.prepare {
        eprintln!("Preparation complete, skipping the rest");
        return Ok(());
    }

    repository
        .populate_upload_directory(&upload_directory, &[hulk_binary])
        .await
        .wrap_err("failed to populate upload directory")?;

    let arguments = &arguments.pre_game;
    let upload_directory = &upload_directory;

    ProgressIndicator::new()
        .map_tasks(
            robots,
            "Executing pregame tasks",
            |robot_address, progress_bar| async move {
                setup_robot(
                    robot_address,
                    upload_directory,
                    arguments,
                    wifi,
                    progress_bar,
                    repository,
                )
                .await
            },
        )
        .await;

    Ok(())
}

async fn setup_robot(
    robot_address: &RobotAddress,
    upload_directory: impl AsRef<Path>,
    arguments: &PreGameArguments,
    wifi: Network,
    progress: ProgressBar,
    repository: &Repository,
) -> Result<()> {
    progress.set_message("Pinging Robot...");
    let robot = Robot::ping_until_available(robot_address.ip).await;

    if !arguments.skip_os_check {
        progress.set_message("Checking OS version...");
        let robot_os_version = robot
            .get_os_version()
            .await
            .wrap_err_with(|| format!("failed to get OS version of {robot_address}"))?;
        let expected_os_version = repository
            .read_os_version()
            .await
            .wrap_err("failed to get configured OS version")?;
        if robot_os_version != expected_os_version {
            bail!(
                "mismatched OS versions: Expected {expected_os_version}, found {robot_os_version}"
            );
        }
    }

    progress.set_message("Stopping HULK...");
    robot
        .execute_systemctl(SystemctlAction::Stop, "hulk")
        .await
        .wrap_err_with(|| format!("failed to stop HULK service on {robot_address}"))?;

    progress.set_message("Uploading: ...");
    robot
        .upload(upload_directory, "hulk", !arguments.no_clean, |status| {
            progress.set_message(format!("Uploading: {status}"))
        })
        .await
        .wrap_err_with(|| format!("failed to upload binary to {robot_address}"))?;

    if wifi != Network::None {
        progress.set_message("Scanning for WiFi...");
        robot
            .scan_networks()
            .await
            .wrap_err_with(|| format!("failed to scan for networks on {robot_address}"))?;
    }

    progress.set_message("Setting WiFi...");
    robot
        .set_wifi(wifi)
        .await
        .wrap_err_with(|| format!("failed to set network on {robot_address}"))?;

    if !arguments.no_restart {
        progress.set_message("Restarting HULK...");
        if let Err(error) = robot
            .execute_systemctl(SystemctlAction::Start, "hulk")
            .await
        {
            let logs = robot
                .retrieve_logs()
                .await
                .wrap_err("failed to retrieve logs")?;
            bail!("failed to restart hulk: {error:#?}\nLogs:\n{logs}")
        };
    }

    Ok(())
}