haze/src/script.rs

use std::{fs::read_to_string, iter::once};

use bollard::Docker;
use camino::Utf8PathBuf;
use miette::{Context, IntoDiagnostic, Report, Result};
use strum::EnumString;

use crate::{
    cloud::CloudOptions,
    config::HazeConfig,
    get_forward_env,
    mapping::{Mapping, MappingSourceType},
    setup,
};

#[derive(Debug, PartialEq, Eq, Clone, Copy, EnumString)]
#[strum(serialize_all = "lowercase")]
pub enum ScriptMode {
    Shell,
    Start,
}

pub async fn run_script(
    path: String,
    config: &HazeConfig,
    docker: &Docker,
    mode_override: Option<ScriptMode>,
) -> Result<()> {
    let path: Utf8PathBuf = Utf8PathBuf::from(path)
        .canonicalize()
        .into_diagnostic()
        .wrap_err("Failed to normalize script path")?
        .try_into()
        .into_diagnostic()
        .wrap_err("Non-utf8 script path")?;
    let contents = read_to_string(&path)
        .into_diagnostic()
        .wrap_err_with(|| format!("Failed to read script file {path}"))?;
    let script = parse_script(&contents, config)
        .wrap_err_with(|| format!("Failed to parse script file {path}"))?;

    let mode = mode_override.unwrap_or(script.mode);
    let mut options = script.options;

    let target_path = Utf8PathBuf::from(format!("/{}", path.file_name().unwrap()));
    let shell = match (script.shell, path.extension()) {
        (Some(shell), _) => shell,
        (None, Some("nu")) => vec!["nu".into()],
        _ => vec!["bash".into()],
    };

    options.mappings.push(
        Mapping::owned(
            MappingSourceType::Absolute,
            path.clone(),
            target_path.clone(),
        )
        .read_only()
        .dont_create(),
    );

    let cloud = setup(docker, options, config, true, true).await?;
    cloud
        .exec(
            docker,
            shell
                .into_iter()
                .chain(once(target_path.into_string()))
                .collect(),
            true,
            get_forward_env(),
        )
        .await?;

    if mode == ScriptMode::Shell {
        cloud.destroy(docker, config).await?;
    }

    Ok(())
}

#[derive(Debug)]
struct Script {
    options: CloudOptions,
    mode: ScriptMode,
    shell: Option<Vec<String>>,
}

fn parse_script(script: &str, config: &HazeConfig) -> Result<Script> {
    let (options, mode, line_num) = script
        .lines()
        .enumerate()
        .find_map(|(i, line)| Some((i, line.strip_prefix("#! haze ")?)))
        .and_then(
            |(i, line)| match (line.strip_prefix("shell"), line.strip_prefix("start")) {
                (Some(shell_options), None) => Some((shell_options.trim(), ScriptMode::Shell, i)),
                (None, Some(script_options)) => Some((script_options.trim(), ScriptMode::Start, i)),
                _ => None,
            },
        )
        .ok_or_else(|| Report::msg("No '#! haze shell' or '#! haze start' found in script"))?;
    let mut options = options.split(' ').peekable();

    let options =
        CloudOptions::parse(config, &mut options).wrap_err("Failed to parse shell options")?;

    let first_content_line = script.lines().nth(line_num + 1).unwrap_or_default();
    let shell = first_content_line.strip_prefix("#!").map(|bang| {
        bang.split(" ")
            .map(str::trim)
            .filter(|s| !s.is_empty())
            .map(String::from)
            .collect()
    });

    Ok(Script {
        options,
        mode,
        shell,
    })
}