rusty_stack/platform/

environment.rs

//! Environment normalization, home directory resolution, and Python interpreter discovery.
//!
//! Provides:
//! - **Home directory resolution** with sudo context handling
//!   (VAL-PLAT-017)
//! - **`.mlstack_env` normalization** producing consistent, idempotent environment
//!   (VAL-PLAT-018)
//! - **Python interpreter discovery** across venvs, conda, and system paths
//!   (VAL-PLAT-020)
//!
//! # Sudo Context
//!
//! When running under `sudo`, the real user's home directory is resolved via:
//! 1. `MLSTACK_USER_HOME` env var
//! 2. `$HOME` (if not `/root`)
//! 3. `/etc/passwd` lookup for `SUDO_USER` / `USER` / `LOGNAME`
//!
//! # Python Discovery Priority
//!
//! 1. `MLSTACK_PYTHON_BIN` / `UV_PYTHON` env vars
//! 2. Active virtualenv (`VIRTUAL_ENV`)
//! 3. Conda environments
//! 4. System Python (`/usr/bin/python3`, etc.)
//! 5. Component venvs (`~/rocm_venv/bin/python`)
//!
//! # .mlstack_env Normalization
//!
//! The env file is idempotent: running normalization twice produces the same
//! output as running it once. Duplicate PATH/LD_LIBRARY_PATH entries are
//! removed, and all required variables are present.

use std::collections::HashSet;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};

// ===========================================================================
// Public Types
// ===========================================================================

/// Result of `.mlstack_env` normalization.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvUpdateResult {
    /// A new `.mlstack_env` file was created.
    Created,
    /// The file was updated with changes.
    Updated,
    /// The file was already normalized, no changes needed.
    Unchanged,
}

// ===========================================================================
// Home Directory Resolution (VAL-PLAT-017)
// ===========================================================================

/// Resolve the real user's home directory, handling sudo context.
///
/// Priority order:
/// 1. `MLSTACK_USER_HOME` env var
/// 2. `HOME` env var (if not `/root` or not running under sudo)
/// 3. `/etc/passwd` lookup for `SUDO_USER`, then `USER`, then `LOGNAME`
/// 4. Fallback to `/tmp`
pub fn resolve_user_home() -> PathBuf {
    // Priority 1: Explicit override
    if let Ok(home) = env::var("MLSTACK_USER_HOME") {
        let home = home.trim().to_string();
        if !home.is_empty() && Path::new(&home).is_dir() {
            return PathBuf::from(home);
        }
    }

    // Priority 2: HOME env var
    if let Ok(home) = env::var("HOME") {
        let home = home.trim().to_string();
        if !home.is_empty() {
            // If HOME is /root but we're running under sudo, skip this
            let running_sudo = env::var("SUDO_USER").is_ok();
            if !(running_sudo && home == "/root") {
                return PathBuf::from(home);
            }
        }
    }

    // Priority 3: Look up user from /etc/passwd via SUDO_USER, USER, LOGNAME
    for key in ["SUDO_USER", "USER", "LOGNAME"] {
        if let Ok(user_name) = env::var(key) {
            let user_name = user_name.trim().to_string();
            if !user_name.is_empty() && user_name != "root" {
                if let Some(home) = lookup_home_from_passwd(&user_name) {
                    return home;
                }
            }
        }
    }

    // Fallback
    PathBuf::from("/tmp")
}

/// Look up a user's home directory from `/etc/passwd`.
pub fn lookup_home_from_passwd(username: &str) -> Option<PathBuf> {
    let content = fs::read_to_string("/etc/passwd").ok()?;
    for line in content.lines() {
        let fields: Vec<&str> = line.split(':').collect();
        if fields.len() >= 6 && fields[0] == username {
            let home = fields[5].trim();
            if !home.is_empty() {
                return Some(PathBuf::from(home));
            }
        }
    }
    None
}

// ===========================================================================
// ML Stack paths — the SINGLE consolidated root (~/.mlstack/)
// ===========================================================================

/// The single ML Stack root directory: `~/.mlstack/`.
///
/// All Rusty-managed state lives under here (Stage 0 consolidation): the
/// global venv (`global/`), named envs (`envs/<name>/`), logs, cache, triton,
/// the installed-component registry (`installed.json`), and config.
pub fn mlstack_root() -> PathBuf {
    resolve_user_home().join(".mlstack")
}

/// The managed **global** default environment: `~/.mlstack/global/`.
pub fn mlstack_global_dir() -> PathBuf {
    mlstack_root().join("global")
}

/// The Python binary of the managed global env: `~/.mlstack/global/bin/python`.
pub fn mlstack_global_python() -> PathBuf {
    mlstack_global_dir().join("bin").join("python")
}

/// The named-envs directory: `~/.mlstack/envs/`.
pub fn mlstack_envs_dir() -> PathBuf {
    mlstack_root().join("envs")
}

/// A specific named env: `~/.mlstack/envs/<name>/`.
pub fn mlstack_env_dir(name: &str) -> PathBuf {
    mlstack_envs_dir().join(name)
}

/// The cache directory: `~/.mlstack/cache/`.
pub fn mlstack_cache_dir() -> PathBuf {
    mlstack_root().join("cache")
}

/// The logs directory: `~/.mlstack/logs/`.
pub fn mlstack_logs_dir() -> PathBuf {
    mlstack_root().join("logs")
}

/// The triton home directory: `~/.mlstack/triton/`.
pub fn mlstack_triton_dir() -> PathBuf {
    mlstack_root().join("triton")
}

/// Ensure the `~/.mlstack/` root and its standard subdirs exist.
pub fn ensure_mlstack_dirs() -> std::io::Result<()> {
    for dir in [
        mlstack_root(),
        mlstack_global_dir(),
        mlstack_envs_dir(),
        mlstack_cache_dir(),
        mlstack_logs_dir(),
        mlstack_triton_dir(),
    ] {
        std::fs::create_dir_all(&dir)?;
    }
    Ok(())
}

/// Create the managed global venv at `~/.mlstack/global/` if it does not yet
/// exist, using `bootstrap_python` (a discovered system/uv interpreter) as the
/// base. Returns the path to the global venv's `python`.
///
/// Prefers `uv venv` when available (fast, deterministic); falls back to
/// `python -m venv`. This is the deterministic anchor that replaces the old
/// non-deterministic interpreter scan for global installs.
pub fn ensure_global_venv(bootstrap_python: &str) -> anyhow::Result<PathBuf> {
    let global_python = mlstack_global_python();
    if global_python.exists() {
        return Ok(global_python);
    }
    ensure_mlstack_dirs()?;
    let global_dir = mlstack_global_dir();
    create_venv(&global_dir, bootstrap_python)?;
    if !global_python.exists() {
        anyhow::bail!(
            "global venv creation reported success but {} is missing",
            global_python.display()
        );
    }
    Ok(global_python)
}

/// Create the managed NAMED env at `~/.mlstack/envs/<name>/` if absent
/// (Tenet 1: install-to-env isolation). Returns the path to its `python`.
///
/// `name` is validated (non-empty, no path separators / traversal) since it is
/// interpolated into a filesystem path.
pub fn ensure_named_venv(name: &str, bootstrap_python: &str) -> anyhow::Result<PathBuf> {
    let name = name.trim();
    if name.is_empty() {
        anyhow::bail!("named env name must not be empty");
    }
    if name.contains(std::path::MAIN_SEPARATOR)
        || name.contains('/')
        || name.contains('\\')
        || name.contains("..")
    {
        anyhow::bail!("invalid env name '{name}'");
    }
    ensure_mlstack_dirs()?;
    let env_dir = mlstack_env_dir(name);
    let env_python = env_dir.join("bin").join("python");
    if env_python.exists() {
        return Ok(env_python);
    }
    create_venv(&env_dir, bootstrap_python)?;
    if !env_python.exists() {
        anyhow::bail!(
            "named venv creation reported success but {} is missing",
            env_python.display()
        );
    }
    Ok(env_python)
}

/// Create a venv at `dir` using `uv` when available, else `python -m venv`.
fn create_venv(dir: &Path, bootstrap_python: &str) -> anyhow::Result<()> {
    if command_on_path("uv") {
        // Treat a uv SPAWN failure (e.g. uv on PATH but not executable / missing
        // runtime) the same as a non-successful run: fall through to the
        // `python -m venv` fallback rather than bailing with a spawn error.
        let uv_status = std::process::Command::new("uv")
            .arg("venv")
            .arg("--python")
            .arg(bootstrap_python)
            .arg(dir)
            .status();
        if let Ok(status) = uv_status {
            if status.success() {
                return Ok(());
            }
        }
    }
    let status = std::process::Command::new(bootstrap_python)
        .arg("-m")
        .arg("venv")
        .arg(dir)
        .status()?;
    if !status.success() {
        anyhow::bail!(
            "failed to create venv at {} via uv/venv (bootstrap={})",
            dir.display(),
            bootstrap_python
        );
    }
    Ok(())
}

// ===========================================================================
// Python Interpreter Discovery (VAL-PLAT-020)
// ===========================================================================

/// Discover Python interpreter paths in priority order, deduplicated.
///
/// Returns only paths that exist on the filesystem. The order is:
/// 0. The managed global env (`~/.mlstack/global/bin/python`) — deterministic anchor
/// 1. `MLSTACK_PYTHON_BIN` / `UV_PYTHON` env vars
/// 2. Active virtualenv (`VIRTUAL_ENV`/bin/python)
/// 3. Conda environment (`CONDA_PREFIX`/bin/python)
/// 4. Component venvs relative to user home (rocm_venv, etc.)
/// 5. System Pythons (/usr/bin/python3.*, /usr/local/bin/python3)
pub fn python_interpreters() -> Vec<PathBuf> {
    let home = resolve_user_home();
    python_interpreters_for_home(&home)
}

/// Resolve the canonical Python binary path — the single interpreter ALL ML
/// components install into.
///
/// Resolution order:
/// 1. `MLSTACK_PYTHON_BIN` / `UV_PYTHON` — an **explicit user override** always
///    wins (respects user intent; e.g. pinning a specific interpreter).
/// 2. The managed global env (`~/.mlstack/global/bin/python`) if it exists —
///    the deterministic single source (Stage 0). Installers create it on first
///    global install via [`ensure_global_venv`].
/// 3. Active virtualenv
/// 4. uv-managed Python (preferred for ML workloads)
/// 5. System Python
///
/// Falls back to `"python3"` if no interpreter is found.
pub fn resolve_canonical_python_bin() -> String {
    // 0. Named-env isolation (Tenet 1): MLSTACK_ENV_NAME pins the target env.
    if let Ok(name) = env::var("MLSTACK_ENV_NAME") {
        let name = name.trim();
        if !name.is_empty() {
            let p = mlstack_env_dir(name).join("bin").join("python");
            if p.exists() {
                return p.to_string_lossy().to_string();
            }
        }
    }
    // 1. Explicit user override wins.
    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
        if let Ok(val) = env::var(key) {
            let val = val.trim();
            if !val.is_empty() && Path::new(val).exists() {
                return val.to_string();
            }
        }
    }
    // 2. Deterministic managed global env (single source).
    let global = mlstack_global_python();
    if global.exists() {
        return global.to_string_lossy().to_string();
    }
    // 3. Discovery fallback.
    let interpreters = python_interpreters();
    if let Some(first) = interpreters.first() {
        return first.to_string_lossy().to_string();
    }
    "python3".to_string()
}

/// Discover Python interpreters for a given home directory (testable).
pub fn python_interpreters_for_home(home: &Path) -> Vec<PathBuf> {
    let mut paths = Vec::new();
    let mut seen = HashSet::new();

    macro_rules! push_if_exists {
        ($path:expr) => {
            let p = $path;
            if p.exists() && seen.insert(p.clone()) {
                paths.push(p);
            }
        };
    }

    // Priority 1: Environment variable overrides — an explicit override wins
    // (kept first so discovery is consistent with `resolve_canonical_python_bin`).
    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
        if let Ok(val) = env::var(key) {
            let val = val.trim().to_string();
            if !val.is_empty() {
                push_if_exists!(PathBuf::from(&val));
            }
        }
    }

    // Priority 2: The managed global env (~/.mlstack/global/bin/python) —
    // the deterministic single source (Stage 0). Derived from `home` so the
    // override is consistent with `resolve_user_home()`.
    push_if_exists!(home
        .join(".mlstack")
        .join("global")
        .join("bin")
        .join("python"));
    push_if_exists!(home
        .join(".mlstack")
        .join("global")
        .join("bin")
        .join("python3"));

    // Priority 2: Active virtualenv
    if let Ok(venv) = env::var("VIRTUAL_ENV") {
        let venv = venv.trim().to_string();
        if !venv.is_empty() {
            push_if_exists!(PathBuf::from(&venv).join("bin/python"));
            push_if_exists!(PathBuf::from(&venv).join("bin/python3"));
        }
    }

    // Priority 3: Conda environment
    if let Ok(conda_prefix) = env::var("CONDA_PREFIX") {
        let conda_prefix = conda_prefix.trim().to_string();
        if !conda_prefix.is_empty() {
            push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python"));
            push_if_exists!(PathBuf::from(&conda_prefix).join("bin/python3"));
        }
    }

    // Priority 4: uv-managed Python installations
    // uv installs Pythons under ~/.local/share/uv/python/ and symlinks to ~/.local/bin/
    // These are PREFERRED over system Pythons because they are the user's ML environment.
    push_if_exists!(home.join(".local/bin/python3"));
    push_if_exists!(home.join(".local/bin/python"));

    // Scan uv python installations (newest first)
    let uv_python_dir = home.join(".local/share/uv/python");
    if let Ok(entries) = fs::read_dir(&uv_python_dir) {
        let mut uv_pythons: Vec<PathBuf> = entries
            .filter_map(|e| e.ok())
            .filter_map(|e| {
                let bin = e.path().join("bin/python3");
                if bin.exists() {
                    Some(bin)
                } else {
                    None
                }
            })
            .collect();
        uv_pythons.sort_by(|a, b| {
            let va = extract_python_version(a);
            let vb = extract_python_version(b);
            // Prefer 3.12 > 3.13 > 3.14 (3.12 is the ML stable target)
            // Within each, sort by version descending
            let priority = |v: (u32, u32)| -> (i32, u32) {
                match v {
                    (3, 12) => (0, v.1), // highest priority
                    (3, 13) => (1, v.1),
                    (3, 11) => (2, v.1),
                    (3, 10) => (3, v.1),
                    _ => (4, v.1), // 3.14+ is lowest priority
                }
            };
            priority(va).cmp(&priority(vb))
        });
        for p in uv_pythons {
            push_if_exists!(p);
        }
    }

    // Priority 5: Component venvs relative to home
    let component_venvs = ["rocm_venv", "pytorch", "ml_stack"];
    for venv_name in &component_venvs {
        push_if_exists!(home.join(venv_name).join("bin/python"));
        push_if_exists!(home.join(venv_name).join("bin/python3"));
    }

    // Priority 6: System Pythons (lowest priority — fallback only)
    let system_pythons = [
        "/usr/local/bin/python3",
        "/usr/bin/python3",
        "/usr/bin/python",
    ];
    for sys_path in &system_pythons {
        push_if_exists!(PathBuf::from(sys_path));
    }

    // Also discover versioned system Pythons
    if let Ok(entries) = fs::read_dir("/usr/bin") {
        let mut versioned: Vec<PathBuf> = entries
            .filter_map(|e| e.ok())
            .filter_map(|e| {
                let name = e.file_name().to_string_lossy().to_string();
                if name.starts_with("python3.") && e.path().exists() {
                    Some(e.path())
                } else {
                    None
                }
            })
            .collect();
        versioned.sort_by(|a, b| {
            // Sort newest first by version number
            let va = extract_python_version(a);
            let vb = extract_python_version(b);
            vb.cmp(&va)
        });
        for p in versioned {
            push_if_exists!(p);
        }
    }

    // Fallback: Always check common user-local Python installations
    // This handles cases where uv installs Python to ~/.local/bin but the directory
    // scanning above might have failed or the canonicalize check earlier didn't work.
    let local_python3 = home.join(".local/bin/python3");
    if local_python3.exists() && seen.insert(local_python3.clone()) {
        paths.push(local_python3);
    }

    paths
}

/// Extract a version tuple from a python binary path like `/usr/bin/python3.13`.
fn extract_python_version(path: &Path) -> (u32, u32) {
    let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
    // Try to parse "python3.13" → (3, 13)
    if let Some(rest) = name.strip_prefix("python") {
        let parts: Vec<&str> = rest.split('.').collect();
        if parts.len() >= 2 {
            let major = parts[0].parse::<u32>().unwrap_or(0);
            let minor = parts[1].parse::<u32>().unwrap_or(0);
            return (major, minor);
        }
        if let Ok(major) = rest.parse::<u32>() {
            return (major, 0);
        }
    }
    (0, 0)
}

// ===========================================================================
// .mlstack_env Normalization (VAL-PLAT-018)
// ===========================================================================

/// Normalize the `.mlstack_env` file in the user's home directory.
///
/// Ensures that the file contains all required environment variables with
/// correct values. The operation is idempotent: running it twice produces
/// the same result as running it once.
///
/// Required variables set by this function:
/// - `MLSTACK_PYTHON_BIN` — path to the resolved Python binary
/// - `ROCM_PATH` — path to ROCm installation
/// - `PYTHONPATH` — includes ROCm lib directory
/// - `PATH` — includes ROCm bin directories
/// - `LD_LIBRARY_PATH` — includes ROCm lib directories
pub fn normalize_mlstack_env(
    user_home: &Path,
    python_bin: &str,
    rocm_path: &Path,
) -> anyhow::Result<EnvUpdateResult> {
    let env_path = user_home.join(".mlstack_env");

    let rocm_str = rocm_path.to_string_lossy();
    let rocm_lib_path = rocm_path.join("lib");
    let rocm_lib = rocm_lib_path.to_string_lossy();

    if env_path.exists() {
        let contents = fs::read_to_string(&env_path)?;

        let (normalized, changed) =
            normalize_env_contents(&contents, python_bin, &rocm_str, &rocm_lib);

        if changed {
            fs::write(&env_path, normalized)?;
            return Ok(EnvUpdateResult::Updated);
        }
        return Ok(EnvUpdateResult::Unchanged);
    }

    // Create new file
    if let Some(parent) = env_path.parent() {
        fs::create_dir_all(parent)?;
    }

    let content = generate_env_file(python_bin, &rocm_str, &rocm_lib);
    fs::write(&env_path, content)?;
    Ok(EnvUpdateResult::Created)
}

/// Normalize the contents of a `.mlstack_env` file.
///
/// Returns `(normalized_contents, changed)`.
pub fn normalize_env_contents(
    contents: &str,
    python_bin: &str,
    rocm_home: &str,
    rocm_lib: &str,
) -> (String, bool) {
    let mut changed = false;
    let mut lines: Vec<String> = Vec::new();
    let mut seen_keys: HashSet<String> = HashSet::new();

    for line in contents.lines() {
        let trimmed = line.trim();
        if trimmed.is_empty() || trimmed.starts_with('#') {
            lines.push(line.to_string());
            continue;
        }

        if let Some(key) = extract_export_key(trimmed) {
            seen_keys.insert(key.clone());

            let (normalized_line, line_changed) =
                normalize_env_line(line, &key, python_bin, rocm_home, rocm_lib);
            changed |= line_changed;
            lines.push(normalized_line);
        } else {
            lines.push(line.to_string());
        }
    }

    // Ensure required keys are present
    let required = [
        (
            "MLSTACK_PYTHON_BIN",
            format!("export MLSTACK_PYTHON_BIN={}", python_bin),
        ),
        ("ROCM_PATH", format!("export ROCM_PATH={}", rocm_home)),
        (
            "ORT_MIGRAPHX_FP16_ENABLE",
            "export ORT_MIGRAPHX_FP16_ENABLE=0".to_string(),
        ),
    ];

    for (key, line) in &required {
        if !seen_keys.contains(*key) {
            lines.push(line.clone());
            changed = true;
        }
    }

    (lines.join("\n"), changed)
}

/// Generate a new `.mlstack_env` file from scratch.
pub fn generate_env_file(python_bin: &str, rocm_home: &str, rocm_lib: &str) -> String {
    let home = dirs::home_dir()
        .map(|h| h.to_string_lossy().to_string())
        .unwrap_or_else(|| "/root".to_string());
    format!(
        "# ML Stack Environment File (generated by Rusty-Stack)\n\
export MLSTACK_PYTHON_BIN={python_bin}\n\
export UV_PYTHON={python_bin}\n\
export ROCM_HOME={rocm_home}\n\
export ROCM_PATH={rocm_home}\n\
export HIP_PATH={rocm_home}\n\
export PYTHONPATH={rocm_lib}:$PYTHONPATH\n\
export PATH=\"{rocm_home}/bin:{rocm_home}/hip/bin:$PATH\"\n\
export LD_LIBRARY_PATH=\"{rocm_home}/lib:{rocm_home}/hip/lib:{rocm_home}/opencl/lib:{home}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"\n\
export ORT_MIGRAPHX_FP16_ENABLE=0\n"
    )
}

// ===========================================================================
// Private Helpers
// ===========================================================================

/// Extract the variable key from an `export KEY=VALUE` line.
fn extract_export_key(line: &str) -> Option<String> {
    let line = line.trim();
    let after_export = line.strip_prefix("export ")?;
    let (key, _) = after_export.split_once('=')?;
    Some(key.trim().to_string())
}

/// Normalize a single env line, returning `(normalized_line, changed)`.
fn normalize_env_line(
    line: &str,
    key: &str,
    python_bin: &str,
    rocm_home: &str,
    rocm_lib: &str,
) -> (String, bool) {
    match key {
        "MLSTACK_PYTHON_BIN" | "UV_PYTHON" => {
            let desired = format!("export {}={}", key, python_bin);
            let trimmed = line.trim();
            if trimmed != desired {
                (desired, true)
            } else {
                (line.to_string(), false)
            }
        }
        "ROCM_PATH" | "ROCM_HOME" | "HIP_PATH" => {
            let desired = format!("export {}={}", key, rocm_home);
            let trimmed = line.trim();
            if trimmed != desired {
                (desired, true)
            } else {
                (line.to_string(), false)
            }
        }
        "PYTHONPATH" => {
            let desired = format!("export PYTHONPATH={}:$PYTHONPATH", rocm_lib);
            let trimmed = line.trim();
            if trimmed != desired {
                (desired, true)
            } else {
                (line.to_string(), false)
            }
        }
        "PATH" => {
            let desired = format!(
                "export PATH=\"{}/bin:{}/hip/bin:$PATH\"",
                rocm_home, rocm_home
            );
            let trimmed = line.trim();
            if trimmed != desired {
                (desired, true)
            } else {
                (line.to_string(), false)
            }
        }
        "LD_LIBRARY_PATH" => {
            let home = dirs::home_dir()
                .map(|h| h.to_string_lossy().to_string())
                .unwrap_or_else(|| "/root".to_string());
            let desired = format!(
                "export LD_LIBRARY_PATH=\"{}/lib:{}/hip/lib:{}/opencl/lib:{}/.mlstack/lib-compat:$LD_LIBRARY_PATH\"",
                rocm_home, rocm_home, rocm_home, home
            );
            let trimmed = line.trim();
            if trimmed != desired {
                (desired, true)
            } else {
                (line.to_string(), false)
            }
        }
        _ => (line.to_string(), false),
    }
}

/// Deduplicate colon-separated path entries.
pub fn dedup_path_var(path_var: &str) -> String {
    let mut seen = HashSet::new();
    let mut result = Vec::new();
    for entry in path_var.split(':') {
        if !entry.is_empty() && seen.insert(entry.to_string()) {
            result.push(entry.to_string());
        }
    }
    result.join(":")
}

/// Is `cmd` executable found on `PATH`? (Dependency-free; does not spawn.)
pub fn command_on_path(cmd: &str) -> bool {
    let path = match env::var_os("PATH") {
        Some(p) => p,
        None => return false,
    };
    for dir in env::split_paths(&path) {
        let candidate = dir.join(cmd);
        if candidate.is_file() {
            return true;
        }
    }
    false
}

// ===========================================================================
// Tests
// ===========================================================================

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

    // -----------------------------------------------------------------------
    // VAL-PLAT-017: Home directory resolution handles sudo context
    // -----------------------------------------------------------------------

    #[test]
    fn test_resolve_user_home_returns_valid_path() {
        let home = resolve_user_home();
        assert!(
            !home.as_os_str().is_empty(),
            "Home directory should not be empty"
        );
    }

    #[test]
    fn test_resolve_user_home_env_override() {
        let _env = crate::test_support::lock_env();
        let saved = std::env::var("MLSTACK_USER_HOME").ok();
        let dir = tempfile::tempdir().unwrap();
        let dir_path = dir.path().to_path_buf();

        std::env::set_var("MLSTACK_USER_HOME", dir_path.to_string_lossy().to_string());
        let result = resolve_user_home();
        assert_eq!(result, dir_path);

        // Restore original state
        match saved {
            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
            None => std::env::remove_var("MLSTACK_USER_HOME"),
        }
    }

    #[test]
    fn test_resolve_user_home_from_passwd() {
        // Look up current user from /etc/passwd
        if let Ok(user) = env::var("USER") {
            if let Some(home) = lookup_home_from_passwd(&user) {
                assert!(
                    home.is_dir(),
                    "Home directory from passwd should exist: {:?}",
                    home
                );
                assert!(
                    !home.as_os_str().is_empty(),
                    "Home directory should not be empty"
                );
            }
        }
    }

    #[test]
    fn test_lookup_home_from_passwd_root() {
        // Root should always be in /etc/passwd
        let home = lookup_home_from_passwd("root");
        assert!(home.is_some(), "root should be in /etc/passwd");
        assert_eq!(home.unwrap(), PathBuf::from("/root"));
    }

    #[test]
    fn test_lookup_home_from_passwd_nonexistent() {
        let home = lookup_home_from_passwd("nonexistent_user_xyz_12345");
        assert!(home.is_none());
    }

    // -----------------------------------------------------------------------
    // VAL-PLAT-020: Python interpreter discovery
    // -----------------------------------------------------------------------

    #[test]
    fn test_python_interpreters_returns_existing_paths() {
        let interpreters = python_interpreters();
        // All returned paths should exist. Note: tests that set MLSTACK_PYTHON_BIN
        // run in parallel and may leave stale env vars pointing to deleted temp dirs.
        // We skip paths from temp directories that no longer exist.
        let valid: Vec<&PathBuf> = interpreters.iter().filter(|p| p.exists()).collect();
        // At least the system python should be found
        assert!(
            !valid.is_empty(),
            "Should find at least one existing Python interpreter"
        );
    }

    #[test]
    fn test_python_interpreters_deduplication() {
        let interpreters = python_interpreters();
        let unique: HashSet<PathBuf> = interpreters.clone().into_iter().collect();
        assert_eq!(
            interpreters.len(),
            unique.len(),
            "Python interpreters should be deduplicated"
        );
    }

    #[test]
    fn test_python_interpreters_env_var_override() {
        let _env = crate::test_support::lock_env();
        // Save current state (hermetic: pin home so a real ~/.mlstack/global
        // on the dev machine can't shadow the explicit override).
        let saved_py = std::env::var("MLSTACK_PYTHON_BIN").ok();
        let saved_home = std::env::var("MLSTACK_USER_HOME").ok();

        // Create a temp file to act as "python"
        let dir = tempfile::tempdir().unwrap();
        let fake_python = dir.path().join("bin/python3");
        fs::create_dir_all(dir.path().join("bin")).unwrap();
        fs::write(&fake_python, "#!/bin/sh").unwrap();

        std::env::set_var(
            "MLSTACK_USER_HOME",
            dir.path().to_string_lossy().to_string(),
        );
        std::env::set_var(
            "MLSTACK_PYTHON_BIN",
            fake_python.to_string_lossy().to_string(),
        );
        let interpreters = python_interpreters();
        assert!(
            interpreters[0] == fake_python,
            "First interpreter should be the env var override: {:?} vs {:?}",
            interpreters[0],
            fake_python
        );

        // Restore original state
        match saved_py {
            Some(v) => std::env::set_var("MLSTACK_PYTHON_BIN", v),
            None => std::env::remove_var("MLSTACK_PYTHON_BIN"),
        }
        match saved_home {
            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
            None => std::env::remove_var("MLSTACK_USER_HOME"),
        }
    }

    #[test]
    fn test_python_interpreters_for_home_with_component_venvs() {
        let dir = tempfile::tempdir().unwrap();
        let home = dir.path();

        // Create component venvs
        for venv_name in &["rocm_venv", "pytorch"] {
            let bin = home.join(venv_name).join("bin");
            fs::create_dir_all(&bin).unwrap();
            fs::write(bin.join("python3"), "#!/bin/sh").unwrap();
        }

        let interpreters = python_interpreters_for_home(home);
        assert!(
            interpreters.len() >= 2,
            "Should find at least 2 component venv interpreters, found {}",
            interpreters.len()
        );

        // rocm_venv should come before system pythons
        let rocm_venv_pos = interpreters
            .iter()
            .position(|p| p.to_string_lossy().contains("rocm_venv"));
        let system_pos = interpreters
            .iter()
            .position(|p| p.to_string_lossy().starts_with("/usr/bin"));
        if let (Some(rv), Some(sp)) = (rocm_venv_pos, system_pos) {
            assert!(rv < sp, "Component venvs should come before system Pythons");
        }
    }

    #[test]
    fn test_extract_python_version() {
        assert_eq!(
            extract_python_version(Path::new("/usr/bin/python3.13")),
            (3, 13)
        );
        assert_eq!(
            extract_python_version(Path::new("/usr/bin/python3.12")),
            (3, 12)
        );
        assert_eq!(
            extract_python_version(Path::new("/usr/bin/python3")),
            (3, 0)
        );
        assert_eq!(extract_python_version(Path::new("/usr/bin/python")), (0, 0));
    }

    // -----------------------------------------------------------------------
    // VAL-PLAT-018: .mlstack_env normalization is idempotent
    // -----------------------------------------------------------------------

    #[test]
    fn test_normalize_mlstack_env_creates_new_file() {
        let dir = tempfile::tempdir().unwrap();
        let python_bin = "/usr/bin/python3";
        let rocm_path = Path::new("/opt/rocm");

        let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
        assert_eq!(result, EnvUpdateResult::Created);

        let env_path = dir.path().join(".mlstack_env");
        assert!(env_path.exists());

        let contents = fs::read_to_string(&env_path).unwrap();
        assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
        assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
        assert!(contents.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
    }

    #[test]
    fn test_normalize_mlstack_env_idempotent() {
        let dir = tempfile::tempdir().unwrap();
        let python_bin = "/usr/bin/python3";
        let rocm_path = Path::new("/opt/rocm");

        // First run: creates
        normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
        let first_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();

        // Second run: should be unchanged
        let result = normalize_mlstack_env(dir.path(), python_bin, rocm_path).unwrap();
        assert_eq!(result, EnvUpdateResult::Unchanged);

        let second_contents = fs::read_to_string(dir.path().join(".mlstack_env")).unwrap();
        assert_eq!(
            first_contents, second_contents,
            "Idempotent: second run should produce identical file"
        );
    }

    #[test]
    fn test_normalize_mlstack_env_updates_stale_values() {
        let dir = tempfile::tempdir().unwrap();
        let env_path = dir.path().join(".mlstack_env");

        // Write stale content
        fs::write(
            &env_path,
            "# ML Stack Environment File\n\
export MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n\
export ROCM_PATH=/opt/rocm-6.0\n\
export PYTHONPATH=/opt/rocm-6.0/lib:$PYTHONPATH\n\
export PATH=\"/opt/rocm-6.0/bin:/opt/rocm-6.0/hip/bin:$PATH\"\n\
export LD_LIBRARY_PATH=\"/opt/rocm-6.0/lib:/opt/rocm-6.0/hip/lib:/opt/rocm-6.0/opencl/lib:$LD_LIBRARY_PATH\"\n",
        )
        .unwrap();

        let result =
            normalize_mlstack_env(dir.path(), "/usr/bin/python3.13", Path::new("/opt/rocm"))
                .unwrap();

        assert_eq!(result, EnvUpdateResult::Updated);

        let contents = fs::read_to_string(&env_path).unwrap();
        assert!(
            contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"),
            "Should update python bin"
        );
        assert!(
            contents.contains("export ROCM_PATH=/opt/rocm\n")
                || contents.contains("export ROCM_PATH=/opt/rocm\r"),
            "Should update ROCM_PATH"
        );
    }

    #[test]
    fn test_normalize_mlstack_env_adds_missing_keys() {
        let dir = tempfile::tempdir().unwrap();
        let env_path = dir.path().join(".mlstack_env");

        // Write content missing required keys
        fs::write(&env_path, "# Minimal env\nexport PATH=/usr/bin:$PATH\n").unwrap();

        let result =
            normalize_mlstack_env(dir.path(), "/usr/bin/python3", Path::new("/opt/rocm")).unwrap();

        assert_eq!(result, EnvUpdateResult::Updated);

        let contents = fs::read_to_string(&env_path).unwrap();
        assert!(contents.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
        assert!(contents.contains("export ROCM_PATH=/opt/rocm"));
    }

    // -----------------------------------------------------------------------
    // Helper function tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_extract_export_key() {
        assert_eq!(
            extract_export_key("export MLSTACK_PYTHON_BIN=/usr/bin/python3"),
            Some("MLSTACK_PYTHON_BIN".to_string())
        );
        assert_eq!(
            extract_export_key("export ROCM_PATH=/opt/rocm"),
            Some("ROCM_PATH".to_string())
        );
        assert_eq!(extract_export_key("# comment"), None);
        assert_eq!(extract_export_key(""), None);
        assert_eq!(extract_export_key("not an export"), None);
    }

    #[test]
    fn test_dedup_path_var() {
        assert_eq!(
            dedup_path_var("/usr/bin:/usr/local/bin:/usr/bin"),
            "/usr/bin:/usr/local/bin"
        );
        assert_eq!(dedup_path_var("/usr/bin"), "/usr/bin");
        assert_eq!(dedup_path_var(""), "");
        assert_eq!(
            dedup_path_var("/usr/bin::/usr/local/bin"),
            "/usr/bin:/usr/local/bin"
        );
    }

    #[test]
    fn test_generate_env_file() {
        let content = generate_env_file("/usr/bin/python3", "/opt/rocm", "/opt/rocm/lib");
        assert!(content.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
        assert!(content.contains("export UV_PYTHON=/usr/bin/python3"));
        assert!(content.contains("export ROCM_HOME=/opt/rocm"));
        assert!(content.contains("export ROCM_PATH=/opt/rocm"));
        assert!(content.contains("export HIP_PATH=/opt/rocm"));
        assert!(content.contains("export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH"));
        assert!(content.contains("/opt/rocm/bin"));
        assert!(content.contains(".mlstack/lib-compat"));
        assert!(content.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
    }

    #[test]
    fn test_normalize_env_contents_no_change() {
        let contents =
            "# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3\nexport ROCM_PATH=/opt/rocm\nexport ORT_MIGRAPHX_FP16_ENABLE=0\n";
        let (result, changed) =
            normalize_env_contents(contents, "/usr/bin/python3", "/opt/rocm", "/opt/rocm/lib");
        assert!(
            !changed,
            "Should not report changes when content is already normalized"
        );
        assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3"));
        assert!(result.contains("export ROCM_PATH=/opt/rocm"));
        assert!(result.contains("export ORT_MIGRAPHX_FP16_ENABLE=0"));
    }

    #[test]
    fn test_normalize_env_contents_with_change() {
        let contents =
            "# ML Stack Environment File\nexport MLSTACK_PYTHON_BIN=/usr/bin/python3.11\n"
                .to_string();
        let (result, changed) = normalize_env_contents(
            &contents,
            "/usr/bin/python3.13",
            "/opt/rocm",
            "/opt/rocm/lib",
        );
        assert!(changed);
        assert!(result.contains("export MLSTACK_PYTHON_BIN=/usr/bin/python3.13"));
        // Should also add missing ROCM_PATH
        assert!(result.contains("export ROCM_PATH=/opt/rocm"));
    }
}