agent-first-psql 0.7.2

use crate::cli::{PsqlAdminAction, PsqlAdminRequest};
use serde::Serialize;
use serde_json::{Value, json};
use std::path::{Path, PathBuf};

const MARKER: &str = "afpsql-managed: true";
const GENERATED_BY: &str = "Generated by afpsql psql install";

pub fn run(req: PsqlAdminRequest) -> i32 {
    let result = handle_action(req.action);
    let stdout = std::io::stdout();
    let mut emitter =
        agent_first_data::CliEmitter::new(stdout.lock(), req.output).with_strict_protocol();
    match result {
        Ok(value) => match emitter.emit_result(value) {
            Ok(()) => 0,
            Err(_) => 4,
        },
        Err(err) => match agent_first_data::json_error("invalid_request", &err.message)
            .hint_if_some(err.hint.as_deref())
            .build()
            .map_err(agent_first_data::CliEmitterError::Build)
        {
            Ok(event) => match emitter.emit(event) {
                Ok(()) => 1,
                Err(_) => 4,
            },
            Err(_) => 4,
        },
    }
}

pub(crate) fn handle_action(action: PsqlAdminAction) -> Result<Value, AdminError> {
    match action {
        PsqlAdminAction::Status { bin_dir } => status(bin_dir),
        PsqlAdminAction::Install { bin_dir } => install(bin_dir),
        PsqlAdminAction::Uninstall { bin_dir } => uninstall(bin_dir),
    }
}

fn status(bin_dir: Option<String>) -> Result<Value, AdminError> {
    let target = resolve_target(bin_dir)?;
    let managed = is_managed_wrapper(&target.wrapper_path).unwrap_or(false);
    let installed = target.wrapper_path.is_file();
    let current_psql = find_on_path(wrapper_file_name())?;
    let active_in_path = current_psql
        .as_ref()
        .map(|p| same_path(p, &target.wrapper_path))
        .unwrap_or(false);

    Ok(admin_result(
        "psql_status",
        json!({
                "wrapper_path": target.wrapper_path,
                "bin_dir": target.bin_dir,
                "installed": installed,
                "managed": managed,
                "active_in_path": active_in_path,
                "current_psql": current_psql,
                "afpsql": current_afpsql_path()?,
        }),
    ))
}

fn install(bin_dir: Option<String>) -> Result<Value, AdminError> {
    let target = resolve_target(bin_dir)?;
    std::fs::create_dir_all(&target.bin_dir).map_err(|e| AdminError::io("create bin dir", e))?;

    if target.wrapper_path.exists() && !is_managed_wrapper(&target.wrapper_path)? {
        return Err(AdminError::invalid_request(
            format!(
                "refusing to overwrite unmanaged psql at {}",
                target.wrapper_path.display()
            ),
            Some("choose another --bin-dir or remove the existing psql yourself".to_string()),
        ));
    }

    let afpsql = current_afpsql_path()?;
    std::fs::write(&target.wrapper_path, wrapper_contents(&afpsql))
        .map_err(|e| AdminError::io("write psql wrapper", e))?;
    set_executable(&target.wrapper_path)?;

    let current_psql = find_on_path(wrapper_file_name())?;
    let active_in_path = current_psql
        .as_ref()
        .map(|p| same_path(p, &target.wrapper_path))
        .unwrap_or(false);

    Ok(admin_result(
        "psql_install",
        json!({
                "wrapper_path": target.wrapper_path,
                "bin_dir": target.bin_dir,
                "installed": true,
                "managed": true,
                "active_in_path": active_in_path,
                "current_psql": current_psql,
                "afpsql": afpsql,
                "hint": if active_in_path { Value::Null } else { json!(format!("add {} to the front of PATH for this psql to take effect", target.bin_dir.display())) },
        }),
    ))
}

fn uninstall(bin_dir: Option<String>) -> Result<Value, AdminError> {
    let target = resolve_target(bin_dir)?;
    if !target.wrapper_path.exists() {
        return Ok(admin_result(
            "psql_uninstall",
            json!({
                    "wrapper_path": target.wrapper_path,
                    "bin_dir": target.bin_dir,
                    "removed": false,
                    "message": "psql wrapper is not installed"
            }),
        ));
    }
    if !is_managed_wrapper(&target.wrapper_path)? {
        return Err(AdminError::invalid_request(
            format!(
                "refusing to remove unmanaged psql at {}",
                target.wrapper_path.display()
            ),
            Some("only wrappers generated by afpsql psql install can be removed".to_string()),
        ));
    }

    std::fs::remove_file(&target.wrapper_path)
        .map_err(|e| AdminError::io("remove psql wrapper", e))?;
    Ok(admin_result(
        "psql_uninstall",
        json!({
                "wrapper_path": target.wrapper_path,
                "bin_dir": target.bin_dir,
                "removed": true
        }),
    ))
}

fn admin_result(code: &str, payload: Value) -> Value {
    let mut payload = payload;
    if let Value::Object(fields) = &mut payload {
        fields.insert("code".to_string(), Value::String(code.to_string()));
    }
    payload
}

struct TargetPath {
    bin_dir: PathBuf,
    wrapper_path: PathBuf,
}

fn resolve_target(bin_dir: Option<String>) -> Result<TargetPath, AdminError> {
    let bin_dir = match bin_dir {
        Some(dir) => expand_tilde(&dir)?,
        None => default_bin_dir()?,
    };
    let wrapper_path = bin_dir.join(wrapper_file_name());
    Ok(TargetPath {
        bin_dir,
        wrapper_path,
    })
}

fn default_bin_dir() -> Result<PathBuf, AdminError> {
    let afpsql = current_afpsql_path()?;
    afpsql.parent().map(Path::to_path_buf).ok_or_else(|| {
        AdminError::invalid_request(
            "cannot determine afpsql executable directory".to_string(),
            Some("pass --bin-dir explicitly".to_string()),
        )
    })
}

fn home_dir() -> Result<PathBuf, AdminError> {
    std::env::var_os("HOME")
        .or_else(|| std::env::var_os("USERPROFILE"))
        .map(PathBuf::from)
        .ok_or_else(|| {
            AdminError::invalid_request(
                "cannot determine home directory".to_string(),
                Some("pass --bin-dir explicitly".to_string()),
            )
        })
}

fn expand_tilde(input: &str) -> Result<PathBuf, AdminError> {
    if input == "~" {
        return home_dir();
    }
    if let Some(rest) = input.strip_prefix("~/") {
        return Ok(home_dir()?.join(rest));
    }
    Ok(PathBuf::from(input))
}

fn wrapper_file_name() -> &'static str {
    if cfg!(windows) { "psql.cmd" } else { "psql" }
}

fn current_afpsql_path() -> Result<PathBuf, AdminError> {
    std::env::current_exe().map_err(|e| AdminError::io("resolve current afpsql path", e))
}

fn wrapper_contents(afpsql: &Path) -> String {
    let afpsql = afpsql.to_string_lossy();
    if cfg!(windows) {
        format!(
            "@echo off\r\nREM {GENERATED_BY}\r\nREM {MARKER}\r\n\"{}\" --mode psql %*\r\n",
            afpsql.replace('"', "\"\"")
        )
    } else {
        format!(
            "#!/usr/bin/env bash\n# {GENERATED_BY}\n# {MARKER}\nexec \"{}\" --mode psql \"$@\"\n",
            shell_double_quote_escape(&afpsql)
        )
    }
}

fn shell_double_quote_escape(value: &str) -> String {
    value
        .replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('$', "\\$")
        .replace('`', "\\`")
}

fn is_managed_wrapper(path: &Path) -> Result<bool, AdminError> {
    if !path.exists() {
        return Ok(false);
    }
    let text = std::fs::read_to_string(path).map_err(|e| AdminError::io("read psql wrapper", e))?;
    Ok(text.contains(MARKER) && text.contains(GENERATED_BY))
}

#[cfg(unix)]
fn set_executable(path: &Path) -> Result<(), AdminError> {
    use std::os::unix::fs::PermissionsExt;
    let mut permissions = std::fs::metadata(path)
        .map_err(|e| AdminError::io("read wrapper metadata", e))?
        .permissions();
    permissions.set_mode(0o755);
    std::fs::set_permissions(path, permissions)
        .map_err(|e| AdminError::io("set wrapper executable bit", e))
}

#[cfg(not(unix))]
fn set_executable(_path: &Path) -> Result<(), AdminError> {
    Ok(())
}

fn find_on_path(name: &str) -> Result<Option<PathBuf>, AdminError> {
    let Some(paths) = std::env::var_os("PATH") else {
        return Ok(None);
    };
    for dir in std::env::split_paths(&paths) {
        for candidate_name in path_candidate_names(name) {
            let candidate = dir.join(candidate_name);
            if candidate.is_file() {
                return Ok(Some(candidate));
            }
        }
    }
    Ok(None)
}

fn path_candidate_names(name: &str) -> Vec<String> {
    if !cfg!(windows) {
        return vec![name.to_string()];
    }
    let lower = name.to_ascii_lowercase();
    if lower.ends_with(".exe") || lower.ends_with(".cmd") || lower.ends_with(".bat") {
        return vec![name.to_string()];
    }
    vec![
        name.to_string(),
        format!("{name}.exe"),
        format!("{name}.cmd"),
        format!("{name}.bat"),
    ]
}

fn same_path(a: &Path, b: &Path) -> bool {
    match (a.canonicalize(), b.canonicalize()) {
        (Ok(a), Ok(b)) => a == b,
        _ => a == b,
    }
}

#[derive(Debug, Serialize)]
pub(crate) struct AdminError {
    message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    hint: Option<String>,
}

impl AdminError {
    fn invalid_request(message: String, hint: Option<String>) -> Self {
        Self { message, hint }
    }

    fn io(action: &str, err: std::io::Error) -> Self {
        Self {
            message: format!("{action} failed: {err}"),
            hint: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn temp_bin_dir(name: &str) -> PathBuf {
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        std::env::temp_dir().join(format!("afpsql_{name}_{}_{}", std::process::id(), suffix))
    }

    #[test]
    fn default_bin_dir_uses_current_exe_parent() {
        let expected = current_afpsql_path()
            .ok()
            .and_then(|p| p.parent().map(Path::to_path_buf));
        assert_eq!(default_bin_dir().ok(), expected);
    }

    #[test]
    fn install_status_uninstall_managed_wrapper() {
        let dir = temp_bin_dir("shim");
        let dir_arg = dir.to_string_lossy().to_string();

        let installed = handle_action(PsqlAdminAction::Install {
            bin_dir: Some(dir_arg.clone()),
        });
        assert!(installed.is_ok());
        let wrapper = dir.join(wrapper_file_name());
        assert!(wrapper.is_file());
        assert!(is_managed_wrapper(&wrapper).unwrap_or(false));

        let status = handle_action(PsqlAdminAction::Status {
            bin_dir: Some(dir_arg.clone()),
        });
        assert!(status.is_ok());
        if let Ok(value) = status {
            assert_eq!(value["installed"], true);
            assert_eq!(value["managed"], true);
        }

        let removed = handle_action(PsqlAdminAction::Uninstall {
            bin_dir: Some(dir_arg),
        });
        assert!(removed.is_ok());
        assert!(!wrapper.exists());
        let _ = std::fs::remove_dir_all(dir);
    }

    #[test]
    fn install_and_uninstall_refuse_unmanaged_psql() {
        let dir = temp_bin_dir("unmanaged");
        assert!(std::fs::create_dir_all(&dir).is_ok());
        let wrapper = dir.join(wrapper_file_name());
        assert!(std::fs::write(&wrapper, "#!/bin/sh\nexit 0\n").is_ok());
        let dir_arg = dir.to_string_lossy().to_string();

        let install = handle_action(PsqlAdminAction::Install {
            bin_dir: Some(dir_arg.clone()),
        });
        assert!(install.is_err());

        let uninstall = handle_action(PsqlAdminAction::Uninstall {
            bin_dir: Some(dir_arg),
        });
        assert!(uninstall.is_err());
        assert!(wrapper.exists());
        let _ = std::fs::remove_dir_all(dir);
    }
}