zerogravity/src/bin/zg.rs

//! `zg` — ZeroGravity daemon manager.
//!
//! All commands exit immediately (safe for agent use via fast-bash MCP).
//! Platform-aware: uses systemd on Linux, launchctl on macOS, and direct
//! process management on Windows.

use std::process::{Command, Stdio};

const SERVICE: &str = "zerogravity";
const PORT: u16 = 8741;

// macOS plist identifier
#[cfg(target_os = "macos")]
const PLIST_LABEL: &str = "com.zerogravity.proxy";

// ANSI colors
const RED: &str = "\x1b[0;31m";
const GREEN: &str = "\x1b[0;32m";
const YELLOW: &str = "\x1b[1;33m";
const CYAN: &str = "\x1b[0;36m";
const BOLD: &str = "\x1b[1m";
const DIM: &str = "\x1b[2m";
const NC: &str = "\x1b[0m";

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let cmd = args.get(1).map(|s| s.as_str()).unwrap_or("");
    let arg = args.get(2).map(|s| s.as_str());

    match cmd {
        "start" => do_start(),
        "stop" => do_stop(),
        "restart" => do_restart(),
        "rebuild" => do_build(),
        "status" => do_status(),
        "logs" => do_logs(arg.unwrap_or("30"), false),
        "logs-follow" => do_logs(arg.unwrap_or("30"), true),
        "logs-all" => do_logs_all(),
        "test" => do_test(arg.unwrap_or("Say hello in exactly 3 words")),
        "health" => do_health(),
        _ => usage(),
    }
}

fn usage() {
    println!("{BOLD}zg{NC} — ZeroGravity daemon manager\n");
    println!("  {CYAN}start{NC}           Start the proxy daemon");
    println!("  {CYAN}stop{NC}            Stop the proxy daemon");
    println!("  {CYAN}restart{NC}         Rebuild + restart");
    println!("  {CYAN}rebuild{NC}         Build release binary only");
    println!("  {CYAN}status{NC}          Service status + quota + usage");
    println!("  {CYAN}logs{NC} [N]        Show last N lines (default 30)");
    println!("  {CYAN}logs-follow{NC} [N] Tail last N lines + follow");
    println!("  {CYAN}logs-all{NC}        Full log dump");
    println!("  {CYAN}test{NC} [msg]      Quick test request (gemini-3-flash)");
    println!("  {CYAN}health{NC}          Health check");
}

fn project_dir() -> String {
    // Resolve project dir from binary location:
    // binary is at <project>/target/release/zg or <project>/target/debug/zg
    let exe = std::env::current_exe().expect("cannot resolve exe path");
    let exe_dir = exe.parent().expect("no parent dir");

    // Walk up from target/release/ or target/debug/
    if let Some(target_dir) = exe_dir.parent() {
        if let Some(project) = target_dir.parent() {
            return project.to_string_lossy().to_string();
        }
    }

    // Fallback: use current dir
    std::env::current_dir()
        .expect("no cwd")
        .to_string_lossy()
        .to_string()
}

#[allow(dead_code)]
fn binary_path() -> String {
    let dir = project_dir();
    #[cfg(windows)]
    {
        format!("{dir}\\target\\release\\zerogravity.exe")
    }
    #[cfg(not(windows))]
    {
        format!("{dir}/target/release/zerogravity")
    }
}

fn base_url() -> String {
    let port = std::env::var("PROXY_PORT")
        .ok()
        .and_then(|p| p.parse::<u16>().ok())
        .unwrap_or(PORT);
    format!("http://localhost:{port}")
}

// ── Platform service management ──

#[cfg(target_os = "linux")]
fn svc_start() -> bool {
    systemctl(&["daemon-reload"]);
    systemctl(&["start", SERVICE])
}

#[cfg(target_os = "linux")]
fn svc_stop() -> bool {
    systemctl(&["stop", SERVICE])
}

#[cfg(target_os = "linux")]
fn svc_status() {
    let output = Command::new("systemctl")
        .args(["--user", "status", SERVICE, "--no-pager"])
        .output();
    match output {
        Ok(o) => {
            let text = String::from_utf8_lossy(&o.stdout);
            for (i, line) in text.lines().enumerate() {
                if i >= 6 {
                    break;
                }
                println!("{line}");
            }
        }
        Err(_) => println!("{RED}Not running{NC}"),
    }
}

#[cfg(target_os = "linux")]
fn svc_logs(n: &str, follow: bool) {
    let mut args = vec!["--user", "-u", SERVICE, "--no-pager", "-n", n];
    if follow {
        args.push("-f");
    }
    let _ = Command::new("journalctl")
        .args(&args)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status();
}

#[cfg(target_os = "linux")]
fn svc_logs_all() {
    let _ = Command::new("journalctl")
        .args(["--user", "-u", SERVICE, "--no-pager"])
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status();
}

#[cfg(target_os = "linux")]
fn systemctl(args: &[&str]) -> bool {
    Command::new("systemctl")
        .arg("--user")
        .args(args)
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

#[cfg(target_os = "linux")]
fn svc_show_fail_logs() {
    let _ = Command::new("journalctl")
        .args(["--user", "-u", SERVICE, "--no-pager", "-n", "20"])
        .status();
}

// ── macOS: launchctl ──

#[cfg(target_os = "macos")]
fn plist_path() -> String {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
    format!("{home}/Library/LaunchAgents/{PLIST_LABEL}.plist")
}

#[cfg(target_os = "macos")]
fn svc_start() -> bool {
    let plist = plist_path();
    if !std::path::Path::new(&plist).exists() {
        eprintln!("{RED}Plist not found at {plist}{NC}");
        eprintln!("  Run ./scripts/setup-macos.sh first");
        return false;
    }
    // bootout first to ensure clean state (ignore errors)
    let uid = Command::new("id")
        .arg("-u")
        .output()
        .ok()
        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
        .unwrap_or_else(|| "501".into());
    let _ = Command::new("launchctl")
        .args(["bootout", &format!("gui/{uid}"), &plist])
        .status();
    Command::new("launchctl")
        .args(["load", &plist])
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

#[cfg(target_os = "macos")]
fn svc_stop() -> bool {
    let plist = plist_path();
    Command::new("launchctl")
        .args(["unload", &plist])
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

#[cfg(target_os = "macos")]
fn svc_status() {
    let output = Command::new("launchctl")
        .args(["list", PLIST_LABEL])
        .output();
    match output {
        Ok(o) if o.status.success() => {
            let text = String::from_utf8_lossy(&o.stdout);
            println!("{text}");
        }
        _ => println!("{RED}Not loaded{NC}"),
    }
}

#[cfg(target_os = "macos")]
fn svc_logs(n: &str, follow: bool) {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
    let log = format!("{home}/Library/Logs/zerogravity.log");
    if !std::path::Path::new(&log).exists() {
        println!("{DIM}(no log file yet){NC}");
        return;
    }
    let mut args = vec!["-n", n];
    if follow {
        args.push("-f");
    }
    args.push(&log);
    let _ = Command::new("tail")
        .args(&args)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status();
}

#[cfg(target_os = "macos")]
fn svc_logs_all() {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
    let log = format!("{home}/Library/Logs/zerogravity.log");
    if std::path::Path::new(&log).exists() {
        let _ = Command::new("cat")
            .arg(&log)
            .stdin(Stdio::inherit())
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .status();
    } else {
        println!("{DIM}(no log file yet){NC}");
    }
}

#[cfg(target_os = "macos")]
fn svc_show_fail_logs() {
    svc_logs("20", false);
}

// ── Windows / other: direct process management ──

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_start() -> bool {
    let bin = binary_path();
    if !std::path::Path::new(&bin).exists() {
        eprintln!("{RED}Binary not found: {bin}{NC}");
        eprintln!("  Run `cargo build --release` first");
        return false;
    }
    // Try starting via scheduled task on Windows
    #[cfg(windows)]
    {
        let result = Command::new("schtasks")
            .args(["/run", "/tn", "ZeroGravity Proxy"])
            .status()
            .map(|s| s.success())
            .unwrap_or(false);
        if result {
            return true;
        }
        // Fallback: start directly (detached)
        eprintln!("{YELLOW}Scheduled task not found, starting directly...{NC}");
    }
    // Fallback for all non-linux/macos: spawn detached
    match Command::new(&bin)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
    {
        Ok(_) => true,
        Err(e) => {
            eprintln!("{RED}Failed to start: {e}{NC}");
            false
        }
    }
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_stop() -> bool {
    #[cfg(windows)]
    {
        let _ = Command::new("schtasks")
            .args(["/end", "/tn", "ZeroGravity Proxy"])
            .status();
    }
    // Also try killing by process name
    #[cfg(windows)]
    {
        let _ = Command::new("taskkill")
            .args(["/F", "/IM", "zerogravity.exe"])
            .status();
    }
    #[cfg(not(windows))]
    {
        let _ = Command::new("pkill").args(["-f", "zerogravity"]).status();
    }
    true
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_status() {
    if health_ok() {
        println!("{GREEN}Running{NC} (responding on port {PORT})");
    } else {
        println!("{RED}Not responding on port {PORT}{NC}");
    }
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_logs(_n: &str, _follow: bool) {
    println!("{DIM}Logs not available via zg on this platform.{NC}");
    println!("  Check the console output where zerogravity was started.");
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_logs_all() {
    svc_logs("0", false);
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn svc_show_fail_logs() {
    println!("{DIM}Check the console output where zerogravity was started.{NC}");
}

// ── Shared helpers ──

fn curl_get(path: &str) -> Option<String> {
    let url = format!("{}{}", base_url(), path);
    Command::new("curl")
        .args(["-sf", &url])
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|o| String::from_utf8_lossy(&o.stdout).to_string())
}

fn curl_post(path: &str, body: &str) -> Option<String> {
    let url = format!("{}{}", base_url(), path);
    Command::new("curl")
        .args([
            "-sf",
            &url,
            "-H",
            "Content-Type: application/json",
            "-d",
            body,
        ])
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|o| String::from_utf8_lossy(&o.stdout).to_string())
}

fn health_ok() -> bool {
    curl_get("/health").is_some()
}

fn jq_print(json: &str) {
    // Try jq first, fall back to raw JSON
    match Command::new("jq")
        .arg(".")
        .stdin(Stdio::piped())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .spawn()
    {
        Ok(mut child) => {
            // Drop stdin before wait so jq sees EOF and doesn't hang
            {
                use std::io::Write;
                if let Some(mut stdin) = child.stdin.take() {
                    let _ = stdin.write_all(json.as_bytes());
                }
            }
            let _ = child.wait();
        }
        Err(_) => {
            // jq not installed — print raw
            println!("{json}");
        }
    }
}

// ── Commands ──

fn do_build() {
    let dir = project_dir();
    println!("{YELLOW}Building release binary...{NC}");
    let status = Command::new("cargo")
        .args(["build", "--release"])
        .current_dir(&dir)
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status()
        .expect("failed to run cargo");

    if status.success() {
        println!("{GREEN}Build complete.{NC}");
    } else {
        eprintln!("{RED}Build failed.{NC}");
        std::process::exit(1);
    }
}

fn do_start() {
    if !svc_start() {
        eprintln!("{RED}Failed to start service.{NC}");
        std::process::exit(1);
    }
    println!("{GREEN}Started.{NC} Waiting for ready...");

    for _ in 0..20 {
        if health_ok() {
            println!("{GREEN}ZeroGravity is up on port {PORT}.{NC}");
            return;
        }
        std::thread::sleep(std::time::Duration::from_millis(500));
    }

    eprintln!("{RED}Proxy didn't become healthy in 10s. Check logs:{NC}");
    svc_show_fail_logs();
    std::process::exit(1);
}

fn do_stop() {
    let _ = svc_stop();
    println!("{YELLOW}Stopped.{NC}");
}

fn do_restart() {
    println!("{YELLOW}Stopping...{NC}");
    do_stop();
    do_build();
    do_start();
}

fn do_status() {
    println!("{BOLD}── Service ──{NC}");
    svc_status();
    println!();

    if !health_ok() {
        println!("{RED}Proxy is not responding on port {PORT}.{NC}");
        std::process::exit(1);
    }

    println!("{BOLD}── Quota ──{NC}");
    match curl_get("/v1/quota") {
        Some(json) => jq_print(&json),
        None => println!("{DIM}(no quota data){NC}"),
    }
    println!();

    println!("{BOLD}── Usage ──{NC}");
    match curl_get("/v1/usage") {
        Some(json) => jq_print(&json),
        None => println!("{DIM}(no usage data){NC}"),
    }
    println!();

    println!("{BOLD}── Sessions ──{NC}");
    match curl_get("/v1/sessions") {
        Some(json) => jq_print(&json),
        None => println!("{DIM}(no sessions){NC}"),
    }
}

fn do_logs(n: &str, follow: bool) {
    svc_logs(n, follow);
}

fn do_logs_all() {
    svc_logs_all();
}

fn do_test(msg: &str) {
    println!("{CYAN}Testing:{NC} {msg}");
    let escaped = msg
        .replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('\n', "\\n")
        .replace('\r', "\\r")
        .replace('\t', "\\t");
    let body =
        format!(r#"{{"model":"gemini-3-flash","input":"{escaped}","stream":false,"timeout":30}}"#);
    match curl_post("/v1/responses", &body) {
        Some(json) => jq_print(&json),
        None => {
            eprintln!("{RED}Request failed. Is the proxy running?{NC}");
            std::process::exit(1);
        }
    }
}

fn do_health() {
    match curl_get("/health") {
        Some(json) => {
            jq_print(&json);
            println!("{GREEN}Healthy{NC}");
        }
        None => {
            println!("{RED}Not responding{NC}");
            std::process::exit(1);
        }
    }
}