jan-cli 0.13.0

//! Package-manager environments for script leaves (`packages.uv`, …).

use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{bail, Context, Result};

use crate::remote::{self, cache_root};
use crate::spec_load;
use crate::{PackagesSpec, RootSpec, RunContext, UvPackages};

/// Result of ensuring a uv env for a leaf.
#[derive(Debug, Clone)]
pub struct UvEnv {
    /// Absolute path to the venv root (`…/packages/uv/<hash>`).
    #[allow(dead_code)]
    pub root: PathBuf,
    /// Absolute path to `bin` (or `Scripts` on Windows).
    pub bin_dir: PathBuf,
}

impl UvEnv {
    #[allow(dead_code)]
    pub fn python3(&self) -> PathBuf {
        if cfg!(windows) {
            self.bin_dir.join("python.exe")
        } else {
            // Prefer python3, fall back to python.
            let p3 = self.bin_dir.join("python3");
            if p3.is_file() {
                p3
            } else {
                self.bin_dir.join("python")
            }
        }
    }
}

/// Merge `packages` along `chain`; deeper nodes override entire per-manager sections.
pub fn collect_chain_packages(chain: &[String], spec: &RootSpec) -> PackagesSpec {
    let mut pkgs = PackagesSpec::default();
    let mut map = &spec.commands;
    for seg in chain {
        let Some(node) = map.get(seg) else { break };
        pkgs.merge_from(node.packages.clone());
        map = &node.commands;
    }
    pkgs
}

/// Ensure package envs for `pkgs` and return an optional uv env to inject.
pub fn ensure_packages(pkgs: &PackagesSpec, ctx: &RunContext<'_>) -> Result<Option<UvEnv>> {
    pkgs.validate(&chain_label(ctx))?;
    if pkgs.pnpm.is_some() {
        bail!("packages.pnpm is not implemented yet");
    }
    let Some(uv) = &pkgs.uv else {
        return Ok(None);
    };
    ensure_uv(uv, ctx).map(Some)
}

fn chain_label(_ctx: &RunContext<'_>) -> String {
    "run".to_string()
}

/// Canonical string used for hashing (stable across process runs).
pub fn uv_canonical(uv: &UvPackages, use_root: &Path) -> Result<String> {
    match uv {
        UvPackages::List(pkgs) => {
            let mut sorted: Vec<String> = pkgs.iter().map(|s| s.trim().to_string()).collect();
            sorted.sort();
            sorted.dedup();
            Ok(format!("list:\n{}", sorted.join("\n")))
        }
        UvPackages::Project(rel) => {
            let dir = spec_load::resolve_under_use_root_any(use_root, rel)?;
            if !dir.is_dir() {
                bail!("packages.uv.project is not a directory: {}", dir.display());
            }
            let pyproject = dir.join("pyproject.toml");
            if !pyproject.is_file() {
                bail!(
                    "packages.uv.project missing pyproject.toml: {}",
                    dir.display()
                );
            }
            let digest = remote::sha256_file(&pyproject)?;
            let lock = dir.join("uv.lock");
            let lock_part = if lock.is_file() {
                format!("\nlock:{}", remote::sha256_file(&lock)?)
            } else {
                String::new()
            };
            Ok(format!(
                "project:{}#{}{lock_part}",
                dir.to_string_lossy(),
                digest
            ))
        }
        UvPackages::Requirements(rel) => {
            let file = spec_load::resolve_under_use_root(use_root, rel)?;
            let digest = remote::sha256_file(&file)?;
            Ok(format!("requirements:{}#{}", file.to_string_lossy(), digest))
        }
    }
}

pub fn uv_hash(canonical: &str) -> String {
    remote::sha256_hex(canonical.as_bytes())[..16].to_string()
}

fn packages_uv_root() -> Result<PathBuf> {
    let d = cache_root()?.join("packages").join("uv");
    fs::create_dir_all(&d).with_context(|| format!("create {}", d.display()))?;
    Ok(d)
}

fn venv_bin_dir(venv: &Path) -> PathBuf {
    if cfg!(windows) {
        venv.join("Scripts")
    } else {
        venv.join("bin")
    }
}

fn uv_available() -> bool {
    Command::new("uv")
        .arg("--version")
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

fn run_uv(args: &[&str], cwd: Option<&Path>) -> Result<()> {
    let mut cmd = Command::new("uv");
    cmd.args(args);
    if let Some(dir) = cwd {
        cmd.current_dir(dir);
    }
    let output = cmd
        .output()
        .with_context(|| format!("spawn `uv {}`", args.join(" ")))?;
    if output.status.success() {
        return Ok(());
    }
    let stderr = String::from_utf8_lossy(&output.stderr);
    let stdout = String::from_utf8_lossy(&output.stdout);
    let detail = if !stderr.trim().is_empty() {
        stderr.trim()
    } else {
        stdout.trim()
    };
    bail!(
        "`uv {}` failed (status {}): {detail}",
        args.join(" "),
        output.status
    );
}

fn ensure_uv(uv: &UvPackages, ctx: &RunContext<'_>) -> Result<UvEnv> {
    if !uv_available() {
        bail!(
            "packages.uv requires `uv` on PATH (install from https://docs.astral.sh/uv/)"
        );
    }
    let use_root = Path::new(&ctx.spec_root.spec_dir);
    let canonical = uv_canonical(uv, use_root)?;
    let hash = uv_hash(&canonical);
    let root = packages_uv_root()?.join(&hash);
    let bin = venv_bin_dir(&root);
    let marker = root.join(".jan-ready");
    let python_ok = bin.join(if cfg!(windows) { "python.exe" } else { "python3" }).is_file()
        || bin.join(if cfg!(windows) { "python.exe" } else { "python" }).is_file();

    if marker.is_file() && python_ok {
        let prev = fs::read_to_string(&marker).unwrap_or_default();
        if prev.trim() == canonical.trim() {
            return Ok(UvEnv {
                root,
                bin_dir: bin,
            });
        }
    }

    fs::create_dir_all(&root).with_context(|| format!("create {}", root.display()))?;

    // Create venv in cache dir.
    if !python_ok {
        run_uv(&["venv", "--clear", root.to_str().unwrap()], None)?;
    }

    let python = if cfg!(windows) {
        bin.join("python.exe")
    } else {
        let p3 = bin.join("python3");
        if p3.is_file() {
            p3
        } else {
            bin.join("python")
        }
    };
    let python_s = python.to_string_lossy();

    match uv {
        UvPackages::List(pkgs) => {
            let req = root.join("requirements.txt");
            let mut sorted: Vec<String> = pkgs.iter().map(|s| s.trim().to_string()).collect();
            sorted.sort();
            sorted.dedup();
            fs::write(&req, format!("{}\n", sorted.join("\n")))
                .with_context(|| format!("write {}", req.display()))?;
            run_uv(
                &[
                    "pip",
                    "install",
                    "--python",
                    python_s.as_ref(),
                    "-r",
                    req.to_str().unwrap(),
                ],
                None,
            )?;
        }
        UvPackages::Requirements(rel) => {
            let file = spec_load::resolve_under_use_root(use_root, rel)?;
            run_uv(
                &[
                    "pip",
                    "install",
                    "--python",
                    python_s.as_ref(),
                    "-r",
                    file.to_str().unwrap(),
                ],
                None,
            )?;
        }
        UvPackages::Project(rel) => {
            let dir = spec_load::resolve_under_use_root_any(use_root, rel)?;
            // Install the project (and its deps) into the cache venv.
            run_uv(
                &[
                    "pip",
                    "install",
                    "--python",
                    python_s.as_ref(),
                    dir.to_str().unwrap(),
                ],
                None,
            )?;
        }
    }

    fs::write(&marker, format!("{canonical}\n"))
        .with_context(|| format!("write {}", marker.display()))?;

    Ok(UvEnv {
        root,
        bin_dir: bin,
    })
}

/// Prepend uv venv bin before other PATH dirs.
pub fn prepend_uv_path(uv: &UvEnv, path_override: Option<String>) -> Result<String> {
    let sep = if cfg!(windows) { ";" } else { ":" };
    let mut parts = vec![uv.bin_dir.to_string_lossy().into_owned()];
    if let Some(rest) = path_override {
        if !rest.is_empty() {
            parts.push(rest);
        }
    } else {
        let current = std::env::var("PATH").unwrap_or_default();
        if !current.is_empty() {
            parts.push(current);
        }
    }
    Ok(parts.join(sep))
}

/// Resolve argv[0]: prefer uv venv for python/python3 when an env is active.
pub fn resolve_program_with_uv(
    program: &str,
    uv: Option<&UvEnv>,
    path_dirs: &[PathBuf],
) -> Result<PathBuf> {
    let trimmed = program.trim();
    let base = Path::new(trimmed)
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or(trimmed);
    if let Some(env) = uv {
        if matches!(base, "python" | "python3" | "python.exe") {
            let candidate = if base.starts_with("python3") && !cfg!(windows) {
                let p3 = env.bin_dir.join("python3");
                if p3.is_file() {
                    p3
                } else {
                    env.bin_dir.join("python")
                }
            } else if cfg!(windows) {
                env.bin_dir.join("python.exe")
            } else {
                env.bin_dir.join(base)
            };
            if candidate.is_file() {
                return Ok(candidate);
            }
        }
    }
    crate::deps::resolve_program(program, path_dirs)
}

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

    #[test]
    fn list_hash_stable_under_reorder() {
        let a = UvPackages::List(vec!["pyyaml".into(), "requests>=2".into()]);
        let b = UvPackages::List(vec!["requests>=2".into(), "pyyaml".into()]);
        let root = Path::new(".");
        let ca = uv_canonical(&a, root).unwrap();
        let cb = uv_canonical(&b, root).unwrap();
        assert_eq!(ca, cb);
        assert_eq!(uv_hash(&ca), uv_hash(&cb));
    }

    #[test]
    fn list_and_requirements_differ() {
        let list = UvPackages::List(vec!["pyyaml".into()]);
        let c = uv_canonical(&list, Path::new(".")).unwrap();
        assert!(c.starts_with("list:"));
        assert_ne!(uv_hash(&c), uv_hash("requirements:/tmp/x#abc"));
    }
}