mecha10-cli 0.8.0

//! Setup command handler
//!
//! Orchestrates environment setup and prerequisites checking.

use crate::context::CliContext;
use crate::ui::style;
use anyhow::Result;
use std::process::Command;

#[derive(Debug, Clone)]
pub struct SetupArgs {
    pub check_only: bool,
    pub install_deps: bool,
    pub quiet: bool,
}

/// Handle setup command
///
/// Checks and installs required dependencies for Mecha10 development.
pub async fn handle_setup(ctx: &mut CliContext, args: &SetupArgs) -> Result<()> {
    println!();
    println!("{}", style::heading("🔧 Mecha10 Environment Setup"));
    println!("{}", style::separator(47));
    println!();

    let mut all_ok = true;

    println!("Checking prerequisites...");
    println!();

    println!("1. Rust toolchain:");
    let rust_ok = binary_available("rustc")?;
    if rust_ok {
        println!("{}", style::success("   ✅ Rust installed"));
    } else {
        println!("{}", style::error("   ❌ Rust not found"));
        all_ok = false;
    }

    println!("2. Cargo:");
    let cargo_ok = binary_available("cargo")?;
    if cargo_ok {
        println!("{}", style::success("   ✅ Cargo installed"));
    } else {
        println!("{}", style::error("   ❌ Cargo not found"));
        all_ok = false;
    }

    println!("3. Docker:");
    let docker = ctx.docker();
    let docker_installed = docker.check_installation().is_ok();
    let docker_daemon = docker.check_daemon().is_ok();
    let docker_ok = docker_installed && docker_daemon;

    if docker_ok {
        println!("{}", style::success("   ✅ Docker installed and running"));
    } else if docker_installed {
        println!("{}", style::warn("   ⚠️  Docker installed but daemon not running"));
        println!("      Start Docker Desktop or run: sudo systemctl start docker");
        all_ok = false;
    } else {
        println!("{}", style::warn("   ⚠️  Docker not available"));
        println!("      Required for control plane services (Redis, PostgreSQL)");
        all_ok = false;
    }

    println!("4. Redis CLI (optional):");
    let redis_cli_ok = binary_available("redis-cli")?;
    if redis_cli_ok {
        println!("{}", style::success("   ✅ redis-cli installed"));
    } else {
        println!("   ⬜ redis-cli not found (optional)");
    }

    println!("5. Python 3:");
    let mut python_ok = check_python()?;
    if python_ok {
        println!("{}", style::success("   ✅ Python 3 installed"));
    } else {
        println!(
            "{}",
            style::warn("   ⚠️  Python 3 not found (needed for AI model quantization)")
        );
        all_ok = false;

        if !args.check_only && args.install_deps {
            println!("   Attempting to install Python 3...");
            let pinned_version = format!("{}.{}", MIN_PYTHON_VERSION.0, MIN_PYTHON_VERSION.1);
            match install_python(&pinned_version).await {
                Ok(true) => {
                    println!("{}", style::success("   ✅ Python 3 installed"));
                    python_ok = check_python()?;
                    if python_ok {
                        all_ok = rust_ok && cargo_ok && docker_ok && python_ok;
                    }
                }
                Ok(false) => {}
                Err(e) => println!("{}", style::error(format!("   ❌ Failed to install Python 3: {}", e))),
            }
        }
    }

    println!();

    if !args.check_only && python_ok {
        setup_venv().await.ok();

        println!("6. Python Dependencies:");
        match install_python_deps(ctx).await {
            Ok(true) => {
                println!("{}", style::success("   ✅ Python dependencies installed in venv"));
                if !args.quiet {
                    println!();
                    println!("   📝 To activate the virtual environment:");
                    println!("      source .venv/bin/activate");
                }
            }
            Ok(false) => println!("   ⬜ No requirements.txt found (skipping)"),
            Err(e) => {
                println!(
                    "{}",
                    style::warn(format!("   ⚠️  Failed to install Python dependencies: {}", e))
                );
                println!("      Install manually with:");
                println!("        source .venv/bin/activate");
                println!("        pip install -r requirements.txt");
            }
        }
        println!();
    }

    if all_ok {
        if !args.quiet {
            println!("✨ Environment setup complete!");
        }
    } else {
        println!("{}", style::warn("⚠️  Some prerequisites are missing"));
        println!();

        if !rust_ok || !cargo_ok {
            println!("Install Rust:");
            println!("  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh");
            println!();
        }

        if !docker_ok {
            println!("Install Docker:");
            println!("  https://docs.docker.com/get-docker/");
            println!();
        }

        if !python_ok {
            println!("Install Python 3:");
            println!(
                "  macOS: brew install python@{}.{}",
                MIN_PYTHON_VERSION.0, MIN_PYTHON_VERSION.1
            );
            println!("  Or: https://www.python.org/downloads/");
            println!();
        }

        if !python_ok && !args.install_deps {
            println!("Run with --install-deps to attempt automatic installation");
        }

        println!();

        if args.check_only {
            return Err(anyhow::anyhow!("Prerequisites check failed"));
        }
    }

    println!();
    Ok(())
}

/// Check whether `<binary> --version` runs successfully, i.e. whether `binary` is installed and
/// on `PATH`. Shared by the rustc/cargo/redis-cli prerequisite checks above - they differ only in
/// which binary name they probe.
fn binary_available(binary: &str) -> Result<bool> {
    let output = Command::new(binary).arg("--version").output();
    Ok(output.map(|o| o.status.success()).unwrap_or(false))
}

const MIN_PYTHON_VERSION: (u32, u32) = (3, 12);

/// Parse a `(major, minor)` version out of `python --version`/`python -V` output,
/// e.g. `"Python 3.9.2\n"` -> `Some((3, 9))`. Returns `None` for unparseable
/// output (unexpected format, non-Python interpreter, etc).
fn parse_python_version(version_output: &str) -> Option<(u32, u32)> {
    let version_part = version_output.trim().strip_prefix("Python ")?;
    let mut parts = version_part.split('.');
    let major = parts.next()?.parse().ok()?;
    let minor = parts.next()?.parse().ok()?;
    Some((major, minor))
}

/// Whether `found` satisfies `minimum` under `(major, minor)` lexicographic
/// comparison (e.g. `(3, 12) >= (3, 9)`, `(3, 9) < (3, 12)`, `(4, 0) >= (3, 12)`).
fn meets_minimum_version(found: (u32, u32), minimum: (u32, u32)) -> bool {
    found >= minimum
}

/// Read the project's pinned Python version from `.python-version` (written by `mecha10 init`).
///
/// An exact `major.minor` pin avoids accepting a `python3` too new for mecha10-rl's PyTorch
/// dependency; falls back to the [`MIN_PYTHON_VERSION`] floor check when no pin file exists.
fn read_pinned_python_version() -> Option<(u32, u32)> {
    read_pinned_python_version_from(std::path::Path::new(".python-version"))
}

/// Same as [`read_pinned_python_version`], parameterized over the file path so tests don't need
/// to mutate the process's current directory (which would race with any other test relying on
/// `std::env::current_dir` under cargo's default parallel test execution).
fn read_pinned_python_version_from(path: &std::path::Path) -> Option<(u32, u32)> {
    let content = std::fs::read_to_string(path).ok()?;
    let mut parts = content.trim().split('.');
    let major = parts.next()?.parse().ok()?;
    let minor = parts.next()?.parse().ok()?;
    Some((major, minor))
}

/// Run `<candidate> --version` and parse the result, tolerating interpreters that print the
/// version to stderr and swallowing spawn failures. Retries once on failure, since a freshly
/// chmod'd executable can briefly lose an exec-visibility race under heavy parallel test load.
fn probe_python_version(candidate: &str) -> Option<(u32, u32)> {
    let output = run_version_probe(candidate)
        .or_else(|| run_version_probe(candidate))
        .or_else(|| run_version_probe(candidate))?;
    parse_python_version(&String::from_utf8_lossy(&output.stdout))
        .or_else(|| parse_python_version(&String::from_utf8_lossy(&output.stderr)))
}

fn run_version_probe(candidate: &str) -> Option<std::process::Output> {
    let output = Command::new(candidate).arg("--version").output().ok()?;
    if !output.status.success() {
        return None;
    }
    Some(output)
}

/// Check if a Python 3 interpreter meeting [`MIN_PYTHON_VERSION`] is installed.
fn check_python() -> Result<bool> {
    Ok(find_python_executable().is_ok())
}

/// Attempt to automatically install a Python 3 interpreter meeting [`MIN_PYTHON_VERSION`],
/// modeled on `scripts/dev/setup.sh`'s `install_python_macos()`. Installs via Homebrew on macOS
/// when available; otherwise prints manual instructions and returns `Ok(false)`.
async fn install_python(pinned_version: &str) -> Result<bool> {
    let brew_available = Command::new("brew")
        .arg("--version")
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false);

    match python_install_strategy(std::env::consts::OS, brew_available) {
        PythonInstallStrategy::Homebrew => {
            println!("   Installing Python {} via Homebrew...", pinned_version);
            let output = tokio::process::Command::new("brew")
                .args(["install", &format!("python@{}", pinned_version)])
                .output()
                .await?;

            if output.status.success() {
                return Ok(true);
            }

            let stderr = String::from_utf8_lossy(&output.stderr);
            println!("   Homebrew install failed: {}", stderr.trim());
            Ok(false)
        }
        PythonInstallStrategy::ManualLinux => {
            println!("Install Python 3:");
            println!("  Linux: apt install python{}", pinned_version);
            println!("  Or: https://www.python.org/downloads/");
            Ok(false)
        }
        PythonInstallStrategy::ManualOther => {
            println!("Install Python 3:");
            println!("  macOS: brew install python@{}", pinned_version);
            println!("  Or: https://www.python.org/downloads/");
            Ok(false)
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PythonInstallStrategy {
    Homebrew,
    ManualLinux,
    ManualOther,
}

fn python_install_strategy(os: &str, brew_available: bool) -> PythonInstallStrategy {
    match os {
        "macos" if brew_available => PythonInstallStrategy::Homebrew,
        "linux" => PythonInstallStrategy::ManualLinux,
        _ => PythonInstallStrategy::ManualOther,
    }
}

/// Setup Python virtual environment
async fn setup_venv() -> Result<bool> {
    let venv_path = std::path::Path::new(".venv");

    if venv_path.exists() {
        return Ok(false);
    }

    let python = find_python_executable()?;

    println!("   Creating virtual environment...");

    let output = tokio::process::Command::new(&python)
        .args(["-m", "venv", ".venv"])
        .output()
        .await?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        anyhow::bail!("venv creation failed: {}", stderr);
    }

    Ok(true)
}

/// Install Python dependencies from requirements.txt using venv
async fn install_python_deps(_ctx: &CliContext) -> Result<bool> {
    let requirements_path = std::path::Path::new("requirements.txt");
    if !requirements_path.exists() {
        return Ok(false);
    }

    let python = if std::path::Path::new(".venv/bin/python").exists() {
        ".venv/bin/python"
    } else if std::path::Path::new(".venv/Scripts/python.exe").exists() {
        ".venv/Scripts/python.exe"
    } else {
        &find_python_executable()?
    };

    println!("   Installing from requirements.txt...");

    let output = tokio::process::Command::new(python)
        .args(["-m", "pip", "install", "-r", "requirements.txt"])
        .env("GIT_TERMINAL_PROMPT", "0")
        .output()
        .await?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        if stderr.contains("could not read Username")
            || stderr.contains("could not read Password")
            || stderr.contains("Permission denied (publickey)")
        {
            anyhow::bail!(
                "pip install failed: {}\n\nfailed to fetch RL training dependencies via git+ssh - \
                ensure your SSH key is added to GitHub (https://github.com/settings/keys) and \
                `ssh -T git@github.com` succeeds.",
                stderr
            );
        }
        anyhow::bail!("pip install failed: {}", stderr);
    }

    Ok(true)
}

/// Find a Python executable on `PATH` that satisfies [`MIN_PYTHON_VERSION`].
///
/// Tries `python3`, `python`, then the pinned-version-specific name, checking the interpreter's
/// actual version (not just that it's Python 3.x) to avoid a `pip install` failing opaquely later.
fn find_python_executable() -> Result<String> {
    let pinned = read_pinned_python_version();

    let versioned = match pinned {
        Some((major, minor)) => format!("python{major}.{minor}"),
        None => format!("python{}.{}", MIN_PYTHON_VERSION.0, MIN_PYTHON_VERSION.1),
    };

    let candidates: [&str; 3] = if pinned.is_some() {
        [&versioned, "python3", "python"]
    } else {
        ["python3", "python", &versioned]
    };

    find_python_executable_among(&candidates, pinned)
}

/// Same as [`find_python_executable`], parameterized over the candidate list and pinned version
/// for testability. With a pin, a candidate must match `major.minor` exactly (PyTorch lacks
/// wheels for every version); without one, falls back to the [`MIN_PYTHON_VERSION`] floor.
fn find_python_executable_among(candidates: &[&str], pinned: Option<(u32, u32)>) -> Result<String> {
    let mut newest_seen: Option<((u32, u32), &str)> = None;

    for &candidate in candidates {
        let Some(version) = probe_python_version(candidate) else {
            continue;
        };

        let qualifies = match pinned {
            Some(pin) => version == pin,
            None => meets_minimum_version(version, MIN_PYTHON_VERSION),
        };
        if qualifies {
            return Ok(candidate.to_string());
        }

        let should_replace = match newest_seen {
            Some((best, _)) => version > best,
            None => true,
        };
        if should_replace {
            newest_seen = Some((version, candidate));
        }
    }

    if let Some(((major, minor), candidate)) = newest_seen {
        if let Some((pin_major, pin_minor)) = pinned {
            anyhow::bail!(
                "Found Python {major}.{minor} at `{candidate}`, but this project pins Python \
                 {pin_major}.{pin_minor} (see .python-version) - PyTorch (a mecha10-rl dependency) \
                 doesn't have wheels for every Python version, so an exact match is required here, \
                 not just \"new enough\". Install Python {pin_major}.{pin_minor} (e.g. `brew install \
                 python@{pin_major}.{pin_minor}` or `pyenv install {pin_major}.{pin_minor}`) and make \
                 sure it's discoverable as `python3`/`python`/`python{pin_major}.{pin_minor}` on PATH, \
                 then re-run `mecha10 setup`.",
            );
        }
        anyhow::bail!(
            "Found Python {major}.{minor} at `{candidate}`, but mecha10 requires Python \
             {req_major}.{req_minor}+ (see mecha10-rl's rl-task-runner/pyproject.toml's \
             `requires-python`). Install a newer Python 3 and make sure it comes first on PATH, \
             then re-run `mecha10 setup`.",
            req_major = MIN_PYTHON_VERSION.0,
            req_minor = MIN_PYTHON_VERSION.1,
        );
    }

    anyhow::bail!("Python 3 not found. Install with: brew install python3 (macOS) or apt install python3 (Linux)")
}

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

    #[test]
    fn binary_available_true_for_a_binary_on_path() {
        assert!(binary_available("cargo").unwrap());
    }

    #[test]
    fn binary_available_false_for_a_nonexistent_binary() {
        assert!(!binary_available("definitely-not-a-real-binary-mecha10-setup-test").unwrap());
    }

    #[test]
    fn parse_python_version_parses_standard_output() {
        assert_eq!(parse_python_version("Python 3.12.1\n"), Some((3, 12)));
        assert_eq!(parse_python_version("Python 3.9.2"), Some((3, 9)));
        assert_eq!(parse_python_version("Python 4.0.0"), Some((4, 0)));
    }

    #[test]
    fn parse_python_version_rejects_unparseable_output() {
        assert_eq!(parse_python_version(""), None);
        assert_eq!(parse_python_version("not python at all"), None);
        assert_eq!(parse_python_version("Python"), None);
        assert_eq!(parse_python_version("Python x.y.z"), None);
    }

    #[test]
    fn meets_minimum_version_compares_major_and_minor() {
        assert!(meets_minimum_version((3, 12), (3, 12)));
        assert!(meets_minimum_version((3, 13), (3, 12)));
        assert!(meets_minimum_version((4, 0), (3, 12)));
        assert!(!meets_minimum_version((3, 11), (3, 12)));
        assert!(!meets_minimum_version((2, 7), (3, 12)));
    }

    /// Builds a tiny shell script that behaves like `<name> --version` printing a fixed Python
    /// version string, for exercising [`find_python_executable_among`] end-to-end without
    /// depending on whatever Python happens to be installed on the test machine.
    #[cfg(unix)]
    fn fake_python_binary(dir: &std::path::Path, name: &str, version_output: &str) -> String {
        use std::fs;
        use std::os::unix::fs::PermissionsExt;

        let path = dir.join(name);
        fs::write(&path, format!("#!/bin/sh\necho \"{version_output}\"\n")).unwrap();
        let mut perms = fs::metadata(&path).unwrap().permissions();
        perms.set_mode(0o755);
        fs::set_permissions(&path, perms).unwrap();
        path.display().to_string()
    }

    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_skips_too_old_interpreter_for_newer_one() {
        let dir = tempfile::tempdir().unwrap();

        let old = fake_python_binary(dir.path(), "old-python", "Python 3.9.2");
        let new = fake_python_binary(dir.path(), "new-python", "Python 3.12.1");

        let found = find_python_executable_among(&[old.as_str(), new.as_str()], None).unwrap();
        assert_eq!(found, new);
    }

    /// Regression test for the Homebrew keg-only bug: `python@3.12` symlinks only the
    /// version-specific binary onto PATH, never the generic `python3`/`python` names. Simulates
    /// that by pointing the generic candidates at nonexistent paths.
    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_falls_back_to_versioned_binary_when_generic_names_missing() {
        let dir = tempfile::tempdir().unwrap();

        let missing_python3 = dir.path().join("definitely-missing-python3").display().to_string();
        let missing_python = dir.path().join("definitely-missing-python").display().to_string();
        let versioned = fake_python_binary(dir.path(), "python3.12", "Python 3.12.4");

        let found = find_python_executable_among(
            &[missing_python3.as_str(), missing_python.as_str(), versioned.as_str()],
            None,
        )
        .unwrap();
        assert_eq!(found, versioned);
    }

    /// A generic `python3` should still be preferred over the versioned fallback name - the
    /// versioned probe exists to find Homebrew keg-only installs, not override a fine interpreter.
    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_prefers_generic_name_over_versioned_fallback() {
        let dir = tempfile::tempdir().unwrap();

        let generic = fake_python_binary(dir.path(), "python3", "Python 3.13.0");
        let versioned = fake_python_binary(dir.path(), "python3.12", "Python 3.12.4");

        let found = find_python_executable_among(&[generic.as_str(), versioned.as_str()], None).unwrap();
        assert_eq!(found, generic);
    }

    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_fails_fast_with_actionable_message_when_only_too_old() {
        let dir = tempfile::tempdir().unwrap();

        let old = fake_python_binary(dir.path(), "old-python", "Python 3.8.10");

        let err = find_python_executable_among(&[old.as_str()], None).unwrap_err();
        let message = err.to_string();
        assert!(
            message.contains("3.8"),
            "message should report the found version: {message}"
        );
        assert!(
            message.contains("3.12"),
            "message should report the required version: {message}"
        );
        assert!(
            message.contains(&old),
            "message should name the interpreter that was found: {message}"
        );
    }

    #[test]
    fn find_python_executable_among_fails_when_nothing_found() {
        let err = find_python_executable_among(&["definitely-not-a-real-interpreter-xyz"], None).unwrap_err();
        assert!(err.to_string().contains("Python 3 not found"));
    }

    /// With a pin, a newer-than-required interpreter must NOT be accepted just because it clears
    /// the old floor - PyTorch (a transitive mecha10-rl dependency) has no wheels for it yet.
    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_rejects_a_newer_interpreter_than_the_pin() {
        let dir = tempfile::tempdir().unwrap();

        let too_new = fake_python_binary(dir.path(), "python3", "Python 3.14.7");

        let err = find_python_executable_among(&[too_new.as_str()], Some((3, 12))).unwrap_err();
        let message = err.to_string();
        assert!(
            message.contains("3.14"),
            "message should report what was found: {message}"
        );
        assert!(
            message.contains("3.12"),
            "message should report the pinned version: {message}"
        );
        assert!(
            message.contains(".python-version"),
            "message should point at the pin's source: {message}"
        );
    }

    /// A pin must also reject an interpreter that's too OLD, same as the floor check would -
    /// pinning isn't only about capping newer versions.
    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_rejects_an_older_interpreter_than_the_pin() {
        let dir = tempfile::tempdir().unwrap();

        let too_old = fake_python_binary(dir.path(), "python3", "Python 3.10.0");

        let err = find_python_executable_among(&[too_old.as_str()], Some((3, 12))).unwrap_err();
        assert!(err.to_string().contains("3.10"));
    }

    /// The whole point of the pin: an interpreter matching it exactly is accepted, even when
    /// a different (unpinned-acceptable) interpreter is also present.
    #[cfg(unix)]
    #[test]
    fn find_python_executable_among_accepts_an_exact_pin_match() {
        let dir = tempfile::tempdir().unwrap();

        let too_new = fake_python_binary(dir.path(), "python3", "Python 3.14.7");
        let pinned_match = fake_python_binary(dir.path(), "python3.12", "Python 3.12.9");

        let found = find_python_executable_among(&[pinned_match.as_str(), too_new.as_str()], Some((3, 12))).unwrap();
        assert_eq!(found, pinned_match);
    }

    #[test]
    fn read_pinned_python_version_parses_major_minor() {
        let dir = tempfile::tempdir().unwrap();
        let version_file = dir.path().join(".python-version");
        std::fs::write(&version_file, "3.12\n").unwrap();

        let result = read_pinned_python_version_from(&version_file);

        assert_eq!(result, Some((3, 12)));
    }

    #[test]
    fn read_pinned_python_version_is_none_when_file_absent() {
        let dir = tempfile::tempdir().unwrap();
        let version_file = dir.path().join(".python-version");

        let result = read_pinned_python_version_from(&version_file);

        assert_eq!(result, None);
    }

    #[test]
    fn python_install_strategy_uses_homebrew_on_macos_when_available() {
        assert_eq!(python_install_strategy("macos", true), PythonInstallStrategy::Homebrew);
    }

    #[test]
    fn python_install_strategy_falls_back_to_manual_on_macos_without_brew() {
        assert_eq!(
            python_install_strategy("macos", false),
            PythonInstallStrategy::ManualOther
        );
    }

    #[test]
    fn python_install_strategy_is_manual_apt_flavored_on_linux() {
        assert_eq!(
            python_install_strategy("linux", true),
            PythonInstallStrategy::ManualLinux
        );
        assert_eq!(
            python_install_strategy("linux", false),
            PythonInstallStrategy::ManualLinux
        );
    }

    #[test]
    fn python_install_strategy_falls_back_to_manual_on_unknown_os() {
        assert_eq!(
            python_install_strategy("windows", true),
            PythonInstallStrategy::ManualOther
        );
        assert_eq!(
            python_install_strategy("freebsd", false),
            PythonInstallStrategy::ManualOther
        );
    }
}