//! The `forge-runner` command line: `register`, `run` and `check`.
//!
//! The standalone `forge-runner` binary and `forge runner` both call [`main_with_args`], so
//! the two behave the same way.

use std::ffi::OsString;
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::Arc;
use std::time::Duration;

use clap::{CommandFactory, Parser, Subcommand};
use forge_ci::proto::RegisterRequest;
use tokio_util::sync::CancellationToken;
use tracing_subscriber::EnvFilter;

use crate::client::ForgeClient;
use crate::config::{self, ExecutorKind, ImagesConfig, Limits, RunnerConfig, UsernsMode};
use crate::executor::{BuiltinExecutor, Executor, PodmanExecutor};
use crate::oci::ImageStore;
use crate::runner::{RunEnd, RunOptions, Runner};
use crate::sandbox::idmap::IdMap;
use crate::update;

/// How often unused images are garbage collected, after the one at startup.
const GC_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);

#[derive(Parser)]
#[command(
    name = "forge-runner",
    version,
    about = "Runs forge CI jobs in a rootless sandbox (builtin) or rootless podman"
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Register this machine with a forge and write the runner config (mode 0600).
    Register {
        /// Forge base URL, for example https://forge.example
        #[arg(long)]
        url: String,
        /// Registration token shown by the forge admin page.
        #[arg(long)]
        token: String,
        /// Runner name shown in the forge.
        #[arg(long)]
        name: String,
        /// Comma-separated runner labels, for example linux,x86_64.
        #[arg(long, value_delimiter = ',')]
        labels: Vec<String>,
        /// Where to write the config. Default: ~/.config/forge-runner/runner.toml
        #[arg(long)]
        config: Option<PathBuf>,
        /// Where job workspaces go. Default: ~/.cache/forge-runner/work
        #[arg(long)]
        work_dir: Option<PathBuf>,
    },
    /// Claim and run jobs until interrupted.
    Run {
        /// The config file. Default: ~/.config/forge-runner/runner.toml
        #[arg(long)]
        config: Option<PathBuf>,
    },
    /// Check that this host can run builtin jobs: user namespaces, overlayfs, cgroup limits,
    /// seccomp and, in subuid mode, subordinate ids. Runs one trivial job through the real sandbox.
    Check {
        /// Where the probe runs. It should be on the filesystem that jobs use.
        /// Default: ~/.cache/forge-runner/work
        #[arg(long)]
        work_dir: Option<PathBuf>,
        /// The runner config, for its `[sandbox] userns` setting. Default: the default config
        /// path, if that file exists, otherwise `auto`.
        #[arg(long)]
        config: Option<PathBuf>,
    },
}

/// Runs the `forge-runner` command line. `args` are the arguments after the program name, so
/// `register --url ...` and not `forge-runner register --url ...`. `bin_name` is shown in help
/// and error messages, for example `forge runner`.
pub fn main_with_args<I, T>(bin_name: &str, args: I) -> ExitCode
where
    I: IntoIterator<Item = T>,
    T: Into<OsString>,
{
    let args: Vec<OsString> = args.into_iter().map(Into::into).collect();
    // The sandbox helper is single-threaded and must not have a tokio runtime, because
    // libcontainer forks and clones. Dispatch it before anything else starts.
    if args.first().is_some_and(is_sandbox_arg) {
        std::process::exit(hidden_dispatch(&args));
    }
    let runtime = match tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
    {
        Ok(runtime) => runtime,
        Err(e) => {
            eprintln!("{bin_name}: cannot start the runtime: {e}");
            return ExitCode::FAILURE;
        }
    };
    // Another subscriber may already be installed, for example by forge-web. Keep that one.
    let _ = tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .with_writer(std::io::stderr)
        .try_init();
    let argv = std::iter::once(OsString::from(bin_name)).chain(args);
    let cli = match Cli::command()
        .bin_name(bin_name)
        .try_get_matches_from(argv)
        .and_then(|matches| <Cli as clap::FromArgMatches>::from_arg_matches(&matches))
    {
        Ok(cli) => cli,
        Err(e) => {
            let code = e.exit_code();
            let _ = e.print();
            return ExitCode::from(code as u8);
        }
    };
    runtime.block_on(async {
        match run(cli).await {
            Ok(code) => code,
            Err(e) => {
                eprintln!("{bin_name}: {e}");
                ExitCode::FAILURE
            }
        }
    })
}

/// Whether this process was started as a hidden sandbox helper: argv[1] is `__sandbox`,
/// `__sandbox-rm` or `__sandbox-unpack`. Callers that have their own runtime or subscriber check
/// this first, and then call [`hidden_main`].
pub fn is_sandbox_invocation() -> bool {
    std::env::args_os()
        .nth(1)
        .is_some_and(|arg| is_sandbox_arg(&arg))
}

/// Runs the hidden helper named by argv[1]. The helpers are single-threaded and must not start
/// a tokio runtime. Returns the process exit code.
pub fn hidden_main() -> i32 {
    let args: Vec<OsString> = std::env::args_os().skip(1).collect();
    hidden_dispatch(&args)
}

fn is_sandbox_arg(arg: &OsString) -> bool {
    arg.as_encoded_bytes().starts_with(b"__sandbox")
}

/// Routes a hidden helper command. `args` start with the command name.
fn hidden_dispatch(args: &[OsString]) -> i32 {
    match args.first().map(|arg| arg.as_encoded_bytes()) {
        Some(b"__sandbox") => crate::sandbox::helper_main(),
        Some(b"__sandbox-rm") => crate::sandbox::idmap::rm_helper_main(&args[1..]),
        Some(b"__sandbox-unpack") => crate::oci::unpack_helper_main(),
        _ => {
            eprintln!("forge-sandbox: unknown helper command");
            125
        }
    }
}

async fn run(cli: Cli) -> Result<ExitCode, Box<dyn std::error::Error>> {
    match cli.command {
        Command::Register {
            url,
            token,
            name,
            labels,
            config: path,
            work_dir,
        } => {
            let client = ForgeClient::new(&url)?;
            let resp = client
                .register(&RegisterRequest {
                    token,
                    name: name.clone(),
                    labels: labels.clone(),
                })
                .await?;
            let path = match path {
                Some(p) => p,
                None => config::default_config_path()?,
            };
            let work_dir = match work_dir {
                Some(w) => w,
                None => config::default_work_dir()?,
            };
            let cfg = RunnerConfig {
                url,
                runner_id: resp.runner_id.clone(),
                runner_secret: resp.runner_secret,
                name,
                labels,
                work_dir,
                concurrency: 1,
                limits: Limits::default(),
                executor: ExecutorKind::default(),
                images: ImagesConfig::default(),
                sandbox: config::SandboxConfig::default(),
                registries: Vec::new(),
                auto_update: true,
            };
            config::save(&path, &cfg)?;
            println!(
                "registered runner {} and wrote {}",
                resp.runner_id,
                path.display()
            );
            Ok(ExitCode::SUCCESS)
        }
        Command::Run { config: path } => {
            let path = match path {
                Some(p) => p,
                None => config::default_config_path()?,
            };
            let cfg = config::load(&path)?;
            let shutdown = CancellationToken::new();
            let executor: Arc<dyn Executor> = match cfg.executor {
                ExecutorKind::Podman => Arc::new(PodmanExecutor::new(cfg.limits.clone())),
                ExecutorKind::Builtin => {
                    // Resolved once. In subuid mode a failure here is a startup error.
                    let map = IdMap::resolve(cfg.sandbox.userns)?;
                    tracing::info!("sandbox ids: {}", map.describe());
                    let helper = std::env::current_exe()?;
                    let root = cfg.work_dir.join("oci");
                    spawn_image_gc(
                        root.clone(),
                        cfg.images.keep_days,
                        map.clone(),
                        helper.clone(),
                        shutdown.clone(),
                    );
                    let images = ImageStore::new(root)
                        .with_max_unpacked_bytes(gib(cfg.limits.image_max_gb))
                        .with_idmap(map.clone(), helper.clone());
                    let jobs_dir = cfg.work_dir.join("jobs");
                    Arc::new(
                        BuiltinExecutor::new(images, cfg.limits.clone(), helper, jobs_dir)
                            .with_idmap(map)
                            .with_registries(&cfg.registries),
                    )
                }
            };
            let runner = Runner::new(cfg, executor, RunOptions::default())?;
            let stop = shutdown.clone();
            tokio::spawn(async move {
                if tokio::signal::ctrl_c().await.is_ok() {
                    tracing::info!("interrupt received, cancelling running jobs");
                    stop.cancel();
                }
            });
            match runner.run(shutdown).await {
                RunEnd::Stopped => Ok(ExitCode::SUCCESS),
                // systemd restarts the runner on this code, and it starts the new binary.
                RunEnd::Updated => Ok(ExitCode::from(update::RESTART_EXIT_CODE)),
            }
        }
        Command::Check {
            work_dir,
            config: config_path,
        } => {
            let work_dir = match work_dir {
                Some(w) => w,
                None => config::default_work_dir()?,
            };
            let mode = check_userns_mode(config_path)?;
            let helper = std::env::current_exe()?;
            if crate::check::run(&helper, &work_dir, mode).await {
                Ok(ExitCode::SUCCESS)
            } else {
                Err("some host checks failed, see the report above".into())
            }
        }
    }
}

/// The `[sandbox] userns` setting for `check`. An explicit config must exist. Without one, the
/// default config is used if it exists, and otherwise the mode is `auto`.
fn check_userns_mode(config_path: Option<PathBuf>) -> Result<UsernsMode, config::ConfigError> {
    if let Some(path) = config_path {
        return Ok(config::load(&path)?.sandbox.userns);
    }
    match config::default_config_path() {
        Ok(path) if path.exists() => Ok(config::load(&path)?.sandbox.userns),
        _ => Ok(UsernsMode::Auto),
    }
}

/// Garbage collects unused images now and then once a day, until `shutdown` is cancelled. Images
/// are removed under `map`, the same as when they were unpacked.
fn spawn_image_gc(
    root: PathBuf,
    keep_days: u32,
    map: IdMap,
    helper: PathBuf,
    shutdown: CancellationToken,
) {
    tokio::spawn(async move {
        loop {
            let (root, map, helper) = (root.clone(), map.clone(), helper.clone());
            let result = tokio::task::spawn_blocking(move || {
                ImageStore::new(root).with_idmap(map, helper).gc(keep_days)
            })
            .await;
            match result {
                Ok(Ok(stats)) => tracing::info!(?stats, "image garbage collection finished"),
                Ok(Err(e)) => tracing::warn!("image garbage collection failed: {e}"),
                Err(e) => tracing::warn!("image garbage collection task failed: {e}"),
            }
            tokio::select! {
                _ = tokio::time::sleep(GC_INTERVAL) => {}
                _ = shutdown.cancelled() => break,
            }
        }
    });
}

/// `image_max_gb` as bytes.
fn gib(gb: u64) -> u64 {
    gb.saturating_mul(1024 * 1024 * 1024)
}