viz-shell 0.1.0

mod banner;
mod build;
mod cli;
mod config;
mod constants;
mod containers;
mod engine;
mod entrypoint;
mod env;
mod mounts;
mod repo;
mod session;
mod share;
mod state;
mod user;

use std::io::IsTerminal;
use std::path::{Path, PathBuf};

use anyhow::{Context, bail, ensure};
use clap::Parser;
use tracing::{debug, info, warn};
use tracing_subscriber::EnvFilter;

use crate::build::ImageStep;
use crate::cli::{Action, Cli};
use crate::config::{Link, RealFiles, Request};
use crate::constants::{
    CONFIG_ENV, CONTAINER_CONFIG_ENV, CONTAINER_ENV, DEFAULT_CONFIG, DEFAULT_LOG_FILTER,
    DEFAULT_STATE_DIR, DOCKER_HOST_ENV, ENTRYPOINT_PATH, GID_ENV, GROUP_ENV, GROUPS_ENV, HOME_ENV,
    HOOKS_ATTACH_ENV, HOOKS_CREATE_ENV, LIBRARY_BASE_FILE, LIBRARY_DIR, LOG_ENV, MOUNTINFO_FILE,
    PASSTHROUGH_ENV, REPO_ENV, SHELL_ENV, SUDO_ENV, UID_ENV, USER_ENV,
};
use crate::containers::{Container, Target};
use crate::engine::{Created, Engine};
use crate::env::CliEnv;
use crate::session::{Enter, Session};
use crate::share::DockerSocket;
use crate::user::User;

fn main() -> anyhow::Result<()> {
    init_tracing();
    let cli = Cli::parse();
    match &cli.action {
        // Before any runtime starts: the entrypoint changes user and execs,
        // which wants a single-threaded process.
        Some(Action::Entrypoint { hold, command }) => entrypoint::run(command, *hold),
        Some(Action::Enter { command }) => entrypoint::enter(command),
        Some(Action::AsUser { command }) => entrypoint::as_user(command),
        Some(Action::Configs) => print_configs(&cli),
        Some(Action::Ls { all }) => runtime()?.block_on(list(*all)),
        Some(Action::Kill { targets, all }) => runtime()?.block_on(kill(targets, *all)),
        Some(Action::New {
            name: Some(name),
            command,
        }) => launched(&cli, Launch::New(name), command),
        Some(Action::New {
            name: None,
            command,
        }) => launched(&cli, Launch::Default, command),
        Some(Action::Attach { target, command }) => {
            let target = target.as_deref().map(Target::parse);
            launched(&cli, Launch::Attach(target), command)
        }
        None if cli.command.is_empty() && !cli.show_effective_config && !cli.show_env => {
            bail!("nothing to run: `vz new` for a shell, `vz -- COMMAND` for a command")
        }
        None => launched(&cli, Launch::Default, &cli.command),
    }
}

fn runtime() -> anyhow::Result<tokio::runtime::Runtime> {
    tokio::runtime::Runtime::new().context("starting the async runtime")
}

/// Exits with the shell's, or the command's, exit code.
fn launched(cli: &Cli, how: Launch, command: &[String]) -> anyhow::Result<()> {
    let exit_code = runtime()?.block_on(launch(cli, how, command))?;
    std::process::exit(exit_code);
}

/// What to load, from the command line, the repository and the host user.
/// A first run writes the library first.
fn request(cli: &Cli, user: &User) -> anyhow::Result<Request> {
    let library_dir = library_dir()?;
    scaffold_library(&library_dir);
    let extra_files = cli
        .files
        .iter()
        .map(|file| {
            std::path::absolute(file).with_context(|| format!("resolving {}", file.display()))
        })
        .collect::<anyhow::Result<_>>()?;
    Ok(Request {
        home: user.home.clone(),
        repo_root: repo::root()?,
        library_dir,
        extra_files,
        config: config_name(cli.config.clone(), std::env::var(CONFIG_ENV).ok())
            .filter(|name| name != DEFAULT_CONFIG),
    })
}

/// The configuration asked for: `-c`, else `VZ_CONFIG`; empty is unset.
/// Not a clap `env`: bare `vz` must list commands whatever the shell exports.
fn config_name(flag: Option<String>, env: Option<String>) -> Option<String> {
    flag.or(env.filter(|name| !name.is_empty()))
}

/// `$XDG_CONFIG_HOME/viz-shell`, or `~/.config/viz-shell` without it.
fn library_dir() -> anyhow::Result<PathBuf> {
    use etcetera::BaseStrategy;
    let strategy = etcetera::choose_base_strategy().context("finding your home folder")?;
    Ok(strategy.config_dir().join(LIBRARY_DIR))
}

/// A first run: the library has no `default`, so vz writes the templates
/// missing from it. A failure leaves the library as it is, never stops vz.
fn scaffold_library(dir: &Path) {
    if !config::needs_scaffold(dir, &RealFiles) {
        return;
    }
    match config::scaffold(dir) {
        Ok(written) => {
            for path in written {
                info!("wrote {}", path.display());
            }
        }
        Err(error) => warn!("no {LIBRARY_BASE_FILE}: {error:#}"),
    }
}

/// Each configuration, the repository's first: its file, what it extends
/// and changes.
fn print_configs(cli: &Cli) -> anyhow::Result<()> {
    let user = User::of_host()?;
    let request = request(cli, &user)?;
    let table = config::read_configs(&request, &RealFiles)?;
    let rows = config::list(&table).into_iter().map(|info| {
        let changes = match info.changes.is_empty() {
            true => "-".to_owned(),
            false => info.changes.join(", "),
        };
        [
            info.name,
            config::shown(&info.file.path, &request.repo_root, &user.home),
            info.extends.unwrap_or_else(|| "-".to_owned()),
            changes,
        ]
    });
    print_table(["NAME", "FILE", "EXTENDS", "CHANGES"], rows);
    Ok(())
}

/// Rows under a header, in columns as wide as their widest cell, like `docker ps`.
fn print_table<const N: usize>(header: [&str; N], rows: impl IntoIterator<Item = [String; N]>) {
    let rows: Vec<[String; N]> = std::iter::once(header.map(str::to_owned))
        .chain(rows)
        .collect();
    let widths: Vec<usize> = (0..N)
        .map(|column| rows.iter().map(|row| row[column].len()).max().unwrap_or(0))
        .collect();
    for row in rows {
        let cells: Vec<String> = row
            .iter()
            .zip(&widths)
            .map(|(cell, width)| format!("{cell:width$}"))
            .collect();
        println!("{}", cells.join("   ").trim_end());
    }
}

/// What a launch does: create a container, or attach to one.
enum Launch<'a> {
    /// `vz new` without a name, or `vz -- COMMAND`: joins a container with
    /// `attach: true` when one fits, else creates one.
    Default,
    /// `vz new NAME`.
    New(&'a str),
    /// `vz attach [TARGET]`.
    Attach(Option<Target>),
}

/// Runs the shell, or `command`, as the host user in a container of the
/// repository: a new one, or one that runs; returns the exit code.
async fn launch(cli: &Cli, how: Launch<'_>, command: &[String]) -> anyhow::Result<i32> {
    ensure!(
        cfg!(target_env = "musl"),
        "vz mounts itself into the container, so it must be a static musl build: \
         cargo build --target x86_64-unknown-linux-musl"
    );
    let user = User::of_host()?;
    debug!("host user: {user:?}");
    let request = request(cli, &user)?;
    let loaded = config::load(&request, &RealFiles).context("in the configuration")?;
    let config_name = request.config.as_deref();
    let config = &loaded.effective;
    let repo_root = request.repo_root.clone();
    // Before anything else: a Dockerfile that requires BASE with no image
    // below it is refused here, engine or not.
    let images =
        build::describe(&config.images, &repo_root, &user.home).context("in the configuration")?;
    if cli.show_effective_config {
        for line in config::header(&loaded.chain, &images, &repo_root, &user.home) {
            println!("{line}");
        }
        print!("{}", config.to_yaml()?);
        return Ok(0);
    }
    debug!("effective configuration: {config:?}");
    let chain = config::links(&loaded.chain, &repo_root, &user.home);
    let config_files: Vec<PathBuf> = loaded
        .chain_files()
        .into_iter()
        .map(|source| source.path)
        .collect();
    // Paths are absolute by now; the repository root is the base of the rest.
    let config_dir = repo_root.clone();

    let cli_env = cli
        .env
        .iter()
        .map(|arg| CliEnv::parse(arg))
        .collect::<anyhow::Result<Vec<_>>>()?;
    let env_paths = env::Paths {
        config_dir: &config_dir,
        repo_root: &repo_root,
        home: &user.home,
    };
    let environment = env::without_reserved(
        env::plan(&config.env, &cli_env, &env_paths)?,
        &reserved_env_names(),
    );
    if cli.show_env {
        print_env(&config::banner_line(&chain), &environment);
        return Ok(0);
    }

    let state_dir = match &config.state_dir {
        Some(dir) => config::resolve_host_path(dir, &config_dir, &user.home),
        None => repo_root.join(DEFAULT_STATE_DIR),
    };
    let state = state::plan(&config.state, &user.home, &state_dir, &repo_root)?;
    let mounts = mounts::plan(&config.mounts, &user.home, &state, &repo_root)?;
    mounts::check_sources_exist(&mounts)?;

    let vz_binary = std::env::current_exe().context("locating the vz binary")?;
    let mount_points = match Path::new(ENTRYPOINT_PATH).exists() {
        true => Some(entrypoint::mount_points(
            &std::fs::read_to_string(MOUNTINFO_FILE).context("reading the mount table")?,
        )),
        false => None,
    };
    session::check_binary_reachable(&vz_binary, mount_points.as_deref())?;

    // Before anything is created on the host: a failed start leaves nothing.
    let engine = Engine::detect().await?;
    let docker = match config.share.docker {
        true => Some(DockerSocket::locate(&engine.docker_endpoint().await?)?),
        false => None,
    };
    let existing = engine.containers(Some(&repo_root)).await?;
    let config_hash = containers::config_hash(&config.to_yaml()?);
    let joining = match &how {
        Launch::Attach(target) => {
            let container = match target {
                Some(target) => containers::find(&existing, target)?,
                None => containers::only_running(&existing)?,
            };
            containers::check_config(container, config_name)?;
            Some(container)
        }
        Launch::Default if config.attach => containers::to_join(&existing, config_name),
        Launch::Default => None,
        Launch::New(name) => {
            containers::check_session_name(name)?;
            if let Ok(taken) = containers::find(&existing, &Target::Name(name.to_string())) {
                bail!("{} exists; attach with `vz attach {name}`", taken.name);
            }
            None
        }
    };

    let workdir = std::env::current_dir().context("reading the current directory")?;
    let tty = std::io::stdin().is_terminal() && std::io::stdout().is_terminal();
    let passthrough = passthrough_env();
    let env_names: Vec<String> = environment.keys().cloned().collect();
    let enter = |container| Enter {
        container,
        workdir: &workdir,
        tty,
        passthrough: &passthrough,
        env_names: &env_names,
        command,
    };
    let show_banner = |container: &str,
                       chain: &[Link],
                       image: &str,
                       image_bases: &[&str],
                       attached: bool,
                       persistent: bool| {
        let Some(art) = &config.banner else {
            return;
        };
        if !(command.is_empty() && tty) {
            return;
        }
        let branch = repo::branch(&repo_root);
        let facts = banner::Session {
            user: &user.name,
            container,
            attached,
            persistent,
            home: &user.home,
            repo_root: &repo_root,
            branch: branch.as_deref(),
            chain,
            config_files: &config_files,
            config: config_name,
            image,
            image_bases,
            shell: config.shell.as_deref(),
            sudo: config.privileges.sudo,
            docker: docker.as_ref().map(|socket| socket.path.as_path()),
            host_network: config.share.host_network,
            mounts: &mounts,
            state_paths: state.len(),
            state_dir: &state_dir,
            env_vars: environment.len(),
            create_hooks: config.hooks.create.len(),
            attach_hooks: config.hooks.attach.len(),
        };
        anstream::print!(
            "{}",
            banner::render(art, &banner::title(&facts), &banner::facts(&facts))
        );
    };
    let values: Vec<(String, String)> = environment
        .iter()
        .map(|(name, var)| (name.clone(), var.value.clone()))
        .collect();

    if let Some(container) = joining {
        if container.config_hash != config_hash {
            warn!(
                "the configuration changed since {} was created; `vz kill {}` and start \
                 again to apply it",
                container.name, container.index
            );
        }
        if !container.running {
            info!("starting {}", container.name);
            engine.start(&container.name).await?;
        }
        // As it was created: its labels hold its chains.
        let created_chain: Vec<Link> = container
            .chain
            .iter()
            .map(|entry| Link::parse(entry))
            .collect();
        let created_bases: Vec<&str> = match container.images.split_last() {
            Some((_, below)) => below.iter().rev().map(String::as_str).collect(),
            None => Vec::new(),
        };
        show_banner(
            &container.name,
            &created_chain,
            &container.image,
            &created_bases,
            true,
            container.persistent,
        );
        return engine.exec(enter(&container.name).args(), &values).await;
    }

    state::create_sources(&state)?;
    mounts::create_points_in_state(&mounts)?;
    let steps = build::plan(&config.image_sources(), &config_dir, &user)?;
    prepare_image(&engine, &steps).await?;
    let (top, below) = steps.split_last().expect("an image chain has a top");
    let image = top.tag().to_owned();
    let image_bases: Vec<&str> = below.iter().rev().map(ImageStep::tag).collect();
    let image_tags: Vec<String> = steps.iter().map(|step| step.tag().to_owned()).collect();
    let chain_entries: Vec<String> = chain.iter().map(Link::entry).collect();
    let session_name = match how {
        Launch::New(name) => Some(name),
        _ => None,
    };
    let repo_dir = repo::dir_name(&repo_root);
    let mut taken = Vec::new();
    let name = loop {
        let index = containers::next_index(&existing, &taken);
        let name = containers::container_name(index, &repo_dir, session_name);
        let labels = containers::Labels {
            repo: &repo_root,
            index,
            session_name,
            config: config_name,
            persistent: config.persistent,
            config_hash: &config_hash,
            chain: &chain_entries,
            images: &image_tags,
        }
        .pairs();
        let session = Session {
            name: &name,
            config: config_name,
            labels: &labels,
            persistent: config.persistent,
            image: &image,
            repo_root: &repo_root,
            workdir: &workdir,
            vz_binary: &vz_binary,
            user: &user,
            state: &state,
            mounts: &mounts,
            docker: docker.as_ref(),
            passthrough: &passthrough,
            env_names: &env_names,
            command,
            tty,
            sudo: config.privileges.sudo,
            host_network: config.share.host_network,
            shell: config.shell.as_deref(),
            hooks: &config.hooks,
        };
        match engine.create(session.create_args(), &values).await? {
            Created::Yes => break name,
            // Another vz took the index meanwhile: the next one.
            Created::NameTaken if taken.len() < MAX_NAME_RETRIES => taken.push(index),
            Created::NameTaken => bail!("{name} is taken, and so are the next names; `vz ls`"),
        }
    };
    if config.persistent {
        engine.start(&name).await?;
        show_banner(&name, &chain, &image, &image_bases, false, true);
        engine.exec(enter(&name).args(), &values).await
    } else {
        show_banner(&name, &chain, &image, &image_bases, false, false);
        engine.start_attached(&name).await
    }
}

/// How often a launch tries the next index when another vz took one.
const MAX_NAME_RETRIES: usize = 10;

/// `vz ls`: this repository's containers, or every repository's.
async fn list(all: bool) -> anyhow::Result<()> {
    let repo_root = match all {
        true => None,
        false => Some(repo::root()?),
    };
    let engine = Engine::detect().await?;
    let found = engine.containers(repo_root.as_deref()).await?;
    if found.is_empty() {
        let whose = if all { "vz" } else { "this repository" };
        println!("No containers of {whose}.");
        return Ok(());
    }
    let row = |container: &Container| {
        [
            container.name.clone(),
            container.config.clone().unwrap_or_else(|| "-".to_owned()),
            (if container.persistent { "yes" } else { "no" }).to_owned(),
            (if container.running {
                "running"
            } else {
                "stopped"
            })
            .to_owned(),
            container.created.clone(),
        ]
    };
    let header = ["NAME", "CONFIG", "PERSISTENT", "STATE", "CREATED"];
    if all {
        let home = User::of_host()?.home;
        let rows = found.iter().map(|container| {
            let [name, config, persistent, state, created] = row(container);
            let repo = config::tilde(&container.repo, &home);
            [name, config, persistent, state, created, repo]
        });
        let [name, config, persistent, state, created] = header;
        print_table([name, config, persistent, state, created, "REPO"], rows);
    } else {
        print_table(header, found.iter().map(row));
    }
    Ok(())
}

/// `vz kill`: removes the named containers of this repository, or all of them.
async fn kill(targets: &[String], all: bool) -> anyhow::Result<()> {
    let repo_root = repo::root()?;
    let engine = Engine::detect().await?;
    let existing = engine.containers(Some(&repo_root)).await?;
    let chosen: Vec<&Container> = match all {
        true => existing.iter().collect(),
        false => targets
            .iter()
            .map(|target| containers::find(&existing, &Target::parse(target)))
            .collect::<anyhow::Result<_>>()?,
    };
    if chosen.is_empty() {
        println!("No containers of this repository.");
        return Ok(());
    }
    let names: Vec<&str> = chosen
        .iter()
        .map(|container| container.name.as_str())
        .collect();
    engine.remove(&names).await?;
    for name in names {
        println!("removed {name}");
    }
    Ok(())
}

/// Names vz sets inside itself: the configured environment cannot change them.
fn reserved_env_names() -> Vec<&'static str> {
    [
        USER_ENV,
        UID_ENV,
        GID_ENV,
        GROUP_ENV,
        HOME_ENV,
        GROUPS_ENV,
        DOCKER_HOST_ENV,
        SUDO_ENV,
        SHELL_ENV,
        CONTAINER_ENV,
        CONTAINER_CONFIG_ENV,
        HOOKS_CREATE_ENV,
        HOOKS_ATTACH_ENV,
        REPO_ENV,
    ]
    .into_iter()
    .chain(PASSTHROUGH_ENV)
    .collect()
}

/// Names and sources, aligned; never a value.
fn print_env(header: &str, environment: &env::Environment) {
    println!("# environment of {header}: names and sources, never values");
    let width = environment.keys().map(String::len).max().unwrap_or(0);
    for (name, var) in environment {
        println!("{name:width$}  {}", var.source);
    }
    println!(
        "# names vz sets itself, which these cannot change: {}",
        reserved_env_names().join(", ")
    );
}

/// Makes the top image of the chain unless the engine has it: builds each
/// image below it the engine lacks, bottom first, then the top; pulls the
/// top when it is a reference. A build pulls what it stands on itself.
async fn prepare_image(engine: &Engine, steps: &[ImageStep]) -> anyhow::Result<()> {
    let (top, below) = steps.split_last().expect("an image chain has a top");
    if engine.has_image(top.tag()).await {
        debug!("image {} is present", top.tag());
        return Ok(());
    }
    for step in below {
        if let ImageStep::Build(plan) = step {
            match engine.has_image(&plan.tag).await {
                true => debug!("image {} is present", plan.tag),
                false => engine.build(plan).await?,
            }
        }
    }
    match top {
        ImageStep::Pull(reference) => engine.pull(reference).await,
        ImageStep::Build(plan) => engine.build(plan).await,
    }
}

fn passthrough_env() -> Vec<(String, String)> {
    PASSTHROUGH_ENV
        .iter()
        .filter_map(|name| Some((name.to_string(), std::env::var(name).ok()?)))
        .collect()
}

/// Events and spans to stderr, filtered by `VZ_LOG`; stdout stays the container's.
fn init_tracing() {
    let filter =
        EnvFilter::try_from_env(LOG_ENV).unwrap_or_else(|_| EnvFilter::new(DEFAULT_LOG_FILTER));
    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_writer(std::io::stderr)
        .with_ansi(std::io::stderr().is_terminal())
        .init();
}

#[cfg(test)]
#[allow(non_snake_case)] // unit__scenario__expected test names
mod tests {
    use super::*;

    #[test]
    fn config_name__flag_and_env__flag_else_nonempty_env() {
        let some = |name: &str| Some(name.to_owned());
        for (flag, env, expected) in [
            (some("ci"), some("sally"), Some("ci")),
            (None, some("sally"), Some("sally")),
            (None, some(""), None),
            (None, None, None),
        ] {
            let case = format!("{flag:?} {env:?}");

            let name = config_name(flag, env);

            assert_eq!(name.as_deref(), expected, "{case}");
        }
    }

    #[test]
    fn reserved_env_names__hooks_and_repo__included() {
        let reserved = reserved_env_names();

        for name in ["VZ_HOOKS_CREATE", "VZ_HOOKS_ATTACH", "VZ_REPO"] {
            assert!(reserved.contains(&name), "{name}: {reserved:?}");
        }
    }
}