aube 1.0.0-beta.2

use super::install::{FrozenMode, InstallOptions};
use clap::{Args, CommandFactory};
use miette::{Context, IntoDiagnostic, miette};

#[derive(Debug, Args)]
// dlx forwards everything after `<command>` to the bin it runs, including
// `--help` and `--version`. Let clap auto-inject its own `-h`/`--help` and
// `--version` handlers and they'd silently swallow those flags before they
// reach the binary — users would see aube's help screen instead of the
// tool's. Disable clap's built-in flags on this subcommand.
//
// `aube dlx --help` on its own (no command) still prints aube's dlx help:
// `params` is optional and the handler intercepts a leading `--help` /
// `-h` before treating anything as a command.
#[command(disable_help_flag = true)]
pub struct DlxArgs {
    /// Command (binary) to run, followed by arguments to pass through to
    /// it.
    ///
    /// The first positional is the command; the rest are forwarded
    /// verbatim to the installed binary. Under `--shell-mode`/`-c` the
    /// positionals are joined and evaluated by `sh -c` instead of
    /// looked up in `node_modules/.bin`.
    #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
    pub params: Vec<String>,
    /// Run the assembled command line through `sh -c` with
    /// `<scratch>/node_modules/.bin` prepended to `PATH`.
    ///
    /// Use this for pipelines, redirects, or env expansion (`aube dlx
    /// -p cowsay -c 'cowsay hello | tr a-z A-Z'`). Mirrors `pnpm dlx
    /// --shell-mode`.
    #[arg(short = 'c', long)]
    pub shell_mode: bool,
    /// Install a specific package (repeatable).
    ///
    /// Overrides inferring from the command.
    #[arg(short = 'p', long = "package")]
    pub package: Vec<String>,
}

/// `aube dlx [-p <pkg>]... <command> [args...]`
///
/// Install one or more packages into a throwaway project and run a binary
/// from them. Matches pnpm's `pnpm dlx` / npm's `npx` surface.
///
/// Flow:
///   1. Create a fresh tempfile::TempDir project with a minimal package.json.
///   2. Run the normal install pipeline there under a CwdGuard that restores
///      the original cwd on drop — including the panic path — so a crash
///      inside install::run can't leave the process with its cwd pointed at
///      an already-removed scratch dir.
///   3. Exec `<tmp>/node_modules/.bin/<command>` from the user's original cwd.
///   4. tempfile removes the scratch dir on drop.
pub async fn run(args: DlxArgs) -> miette::Result<()> {
    let DlxArgs {
        params,
        package,
        shell_mode,
    } = args;

    // Bare `aube dlx` or `aube dlx --help` / `-h` prints aube's dlx help.
    // Once a command is present, any further flags (including `--help`)
    // belong to the installed binary.
    let first = params.first().map(String::as_str);
    if matches!(first, None | Some("--help" | "-h")) && package.is_empty() {
        crate::Cli::command()
            .find_subcommand_mut("dlx")
            .expect("dlx is a registered subcommand")
            .print_help()
            .map_err(|e| miette!("failed to render help: {e}"))?;
        println!();
        return Ok(());
    }

    // When only `-p` is given, dlx needs at least one arg to serve as the
    // bin name; without it, we don't know which binary to exec.
    let command = params
        .first()
        .cloned()
        .ok_or_else(|| miette!("dlx: missing command to run"))?;
    let bin_args: Vec<String> = params.iter().skip(1).cloned().collect();

    // Derive the packages to install. `-p` wins; otherwise the command name
    // is the package name (the common `pnpm dlx <pkg>` case). Under
    // `--shell-mode` the first positional is a shell line, not a bin name,
    // so we fall back to the first whitespace-separated word for inference
    // when `-p` wasn't given — same as pnpm.
    let install_specs: Vec<String> = if package.is_empty() {
        if shell_mode {
            let first_word = command
                .split_whitespace()
                .next()
                .ok_or_else(|| miette!("dlx --shell-mode: missing command line to run"))?;
            vec![first_word.to_string()]
        } else {
            vec![command.clone()]
        }
    } else {
        package
    };

    // Bin name is only used in the non-shell path. Under shell-mode the
    // user assembles their own line and we run it through `sh -c`, so any
    // bin lookup is the shell's job.
    let bin_name = bin_name_for(&command);

    let tmp = tempfile::Builder::new()
        .prefix("aube-dlx-")
        .tempdir()
        .into_diagnostic()
        .wrap_err("failed to create dlx scratch dir")?;
    let project_dir = tmp.path().to_path_buf();

    // Minimal package.json. Version specs and dist-tags pass through as-is
    // — the resolver handles them exactly as it would from a real manifest.
    let deps: serde_json::Map<String, serde_json::Value> = install_specs
        .iter()
        .map(|spec| {
            let (name, version) = split_spec(spec);
            (
                name.to_string(),
                serde_json::Value::String(version.to_string()),
            )
        })
        .collect();
    let manifest = serde_json::json!({
        "name": "aube-dlx",
        "version": "0.0.0",
        "private": true,
        "dependencies": deps,
    });
    std::fs::write(
        project_dir.join("package.json"),
        serde_json::to_string_pretty(&manifest).into_diagnostic()?,
    )
    .into_diagnostic()
    .wrap_err("failed to write dlx package.json")?;

    // install::run pulls its project dir from std::env::current_dir(), which
    // is process-global state. The CwdGuard below captures the current dir,
    // switches into the scratch project for the duration of the install, and
    // restores the original on drop — so the exec path below, any error
    // diagnostic rendering, and even a panic unwinding past this frame all
    // observe the user's real cwd instead of a dir that's about to vanish.
    let prev_cwd = {
        let _cwd_guard = CwdGuard::switch_to(&project_dir)?;
        let install_result = super::install::run(InstallOptions::with_mode(FrozenMode::No)).await;
        let prev = _cwd_guard.original.clone();
        install_result.wrap_err("dlx install failed")?;
        prev
        // _cwd_guard drops here, restoring cwd.
    };

    // Run from the user's original cwd so the invoked tool sees their
    // project, not the scratch dir — this matches pnpm dlx.
    //
    // Under `--shell-mode` we evaluate the joined positionals via `sh -c`
    // with the scratch project's `node_modules/.bin` prepended to PATH,
    // so pipelines/redirects work and the freshly installed bin
    // resolves first. Otherwise we exec the bin directly so its argv
    // round-trips bit-for-bit.
    let status = if shell_mode {
        let line = std::iter::once(command.as_str())
            .chain(bin_args.iter().map(String::as_str))
            .collect::<Vec<_>>()
            .join(" ");
        // `dlx` installs into a scratch tempdir, which honors `modulesDir`
        // from the user's `~/.npmrc` / `aube-workspace.yaml` if it's set
        // globally. Read the same setting here so the scratch bin dir
        // matches where the install actually wrote the bins.
        let bin_dir = super::project_modules_dir(&project_dir).join(".bin");
        let new_path = aube_scripts::prepend_path(&bin_dir);
        let mut cmd = aube_scripts::spawn_shell(&line);
        cmd.env("PATH", &new_path)
            .current_dir(&prev_cwd)
            .stderr(aube_scripts::child_stderr())
            .status()
            .await
            .into_diagnostic()
            .wrap_err("failed to execute dlx shell command")?
    } else {
        let bin_path = super::project_modules_dir(&project_dir)
            .join(".bin")
            .join(&bin_name);
        if !bin_path.exists() {
            return Err(miette!(
                "dlx: binary not found after install: {bin_name}\n\
                 help: the package may ship the binary under a different name — try `aube dlx -p <package> <bin>`"
            ));
        }
        tokio::process::Command::new(&bin_path)
            .args(&bin_args)
            .current_dir(&prev_cwd)
            .stderr(aube_scripts::child_stderr())
            .status()
            .await
            .into_diagnostic()
            .wrap_err("failed to execute dlx binary")?
    };

    // tmp drops here, removing the scratch project.
    drop(tmp);

    if !status.success() {
        std::process::exit(status.code().unwrap_or(1));
    }
    Ok(())
}

/// RAII guard that swaps the process cwd on construction and restores it
/// on drop — including when the enclosing scope unwinds due to a panic.
struct CwdGuard {
    original: std::path::PathBuf,
}

impl CwdGuard {
    fn switch_to(new_dir: &std::path::Path) -> miette::Result<Self> {
        let original = std::env::current_dir().into_diagnostic()?;
        std::env::set_current_dir(new_dir)
            .into_diagnostic()
            .wrap_err("failed to switch into dlx scratch dir")?;
        Ok(Self { original })
    }
}

impl Drop for CwdGuard {
    fn drop(&mut self) {
        // Best-effort: if restoring the cwd fails we can't meaningfully
        // recover, but we also don't want to double-panic from Drop.
        let _ = std::env::set_current_dir(&self.original);
    }
}

/// Strip any `@version` suffix from a package spec, preserving `@scope/`
/// prefixes. Dlx defaults to the `latest` dist-tag when no version is given,
/// so the spec arm of `split_name_spec` becomes `"latest"` here.
fn split_spec(spec: &str) -> (&str, &str) {
    let (name, version) = super::split_name_spec(spec);
    (name, version.unwrap_or("latest"))
}

/// The binary name `aube dlx <cmd>` should resolve to — strip any version
/// suffix and any `@scope/` prefix, since `node_modules/.bin/` is flat and
/// scoped packages still land under their unscoped bin name.
fn bin_name_for(command: &str) -> String {
    let (name, _) = split_spec(command);
    name.rsplit('/').next().unwrap_or(name).to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn split_spec_plain() {
        assert_eq!(split_spec("cowsay"), ("cowsay", "latest"));
    }

    #[test]
    fn split_spec_versioned() {
        assert_eq!(split_spec("cowsay@1.5.0"), ("cowsay", "1.5.0"));
    }

    #[test]
    fn split_spec_scoped() {
        assert_eq!(split_spec("@scope/foo"), ("@scope/foo", "latest"));
    }

    #[test]
    fn split_spec_scoped_versioned() {
        assert_eq!(split_spec("@scope/foo@2.0.0"), ("@scope/foo", "2.0.0"));
    }

    #[test]
    fn bin_name_strips_scope_and_version() {
        assert_eq!(bin_name_for("cowsay@1.5.0"), "cowsay");
        assert_eq!(bin_name_for("@scope/foo@2"), "foo");
        assert_eq!(bin_name_for("@scope/foo"), "foo");
    }
}