workmux 0.1.213

//! OMP agent status tracking setup.
//!
//! Detects OMP via its agent directory at `~/.omp/agent/` and writes
//! `workmux-status.ts` to that agent directory.
//!
//! Installs extension by writing `workmux-status.ts` to the extensions directory.

use anyhow::{Context, Result};
use std::ffi::OsString;
use std::fs;
use std::path::{Path, PathBuf};

use super::StatusCheck;

/// The OMP extension source, embedded at compile time.
const EXTENSION_SOURCE: &str = include_str!("../../.omp/extensions/workmux-status.ts");

fn omp_agent_dir() -> Option<PathBuf> {
    let home = home::home_dir()?;
    Some(omp_agent_dir_with_env(&home, |key| std::env::var_os(key)))
}

pub(crate) fn omp_agent_dir_with_env(
    home: &Path,
    get_env: impl Fn(&str) -> Option<OsString>,
) -> PathBuf {
    if let Some(dir) = get_env("PI_CODING_AGENT_DIR") {
        return PathBuf::from(dir);
    }

    let config_dir = get_env("PI_CONFIG_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from(".omp"));

    if config_dir.is_absolute() {
        config_dir.join("agent")
    } else {
        home.join(config_dir).join("agent")
    }
}

fn extension_path() -> Option<PathBuf> {
    omp_agent_dir().map(|d| d.join("extensions/workmux-status.ts"))
}

/// Detect if OMP is present via filesystem.
/// Returns the reason string if detected, None otherwise.
pub fn detect() -> Option<&'static str> {
    if omp_agent_dir().is_some_and(|d| d.is_dir()) {
        return Some("found ~/.omp/agent/");
    }
    None
}

/// Check if workmux extension is installed for OMP.
pub fn check() -> Result<StatusCheck> {
    let Some(path) = extension_path() else {
        return Ok(StatusCheck::NotInstalled);
    };

    if path.exists() {
        Ok(StatusCheck::Installed)
    } else {
        Ok(StatusCheck::NotInstalled)
    }
}

/// Install workmux extension for OMP.
/// Returns a description of what was done.
pub fn install() -> Result<String> {
    let path =
        extension_path().ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;

    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).context("Failed to create omp extensions directory")?;
    }

    fs::write(&path, EXTENSION_SOURCE).context("Failed to write omp extension")?;

    Ok(format!(
        "Installed extension to {}. Restart omp for it to take effect.",
        path.display()
    ))
}

/// Remove workmux extension for OMP agent.
///
/// Deletes the extension file and cleans up empty parent directories.
pub fn uninstall() -> Result<String> {
    let Some(path) = extension_path() else {
        return Ok("omp config dir not found, nothing to uninstall".to_string());
    };
    uninstall_at(path)
}

fn uninstall_at(path: PathBuf) -> Result<String> {
    if !path.exists() {
        return Ok("No omp extension found".to_string());
    }
    fs::remove_file(&path)?;
    // Clean up empty extensions directory
    if let Some(parent) = path.parent()
        && parent.read_dir().is_ok_and(|mut it| it.next().is_none())
    {
        let _ = fs::remove_dir(parent);
    }
    Ok(format!("Removed omp extension at {}", path.display()))
}

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

    #[test]
    fn test_extension_tracks_waiting_status() {
        assert!(EXTENSION_SOURCE.contains("lastStatus"));
        assert!(EXTENSION_SOURCE.contains("statusQueue"));
        assert!(EXTENSION_SOURCE.contains("status === lastStatus"));
        assert!(EXTENSION_SOURCE.contains("pi.on(\"message_end\""));
        assert!(EXTENSION_SOURCE.contains("\"role\" in event.message"));
        assert!(EXTENSION_SOURCE.contains("event.message.role === \"assistant\""));
        assert!(EXTENSION_SOURCE.contains("pi.on(\"tool_call\""));
        assert!(EXTENSION_SOURCE.contains("event.toolName === \"ask\""));
        assert!(EXTENSION_SOURCE.contains("setStatus(\"waiting\")"));
    }

    #[test]
    fn test_omp_agent_dir_default() {
        let dir = omp_agent_dir_with_env(Path::new("/home/test"), |_| None);

        assert_eq!(dir, PathBuf::from("/home/test/.omp/agent"));
    }

    #[test]
    fn test_omp_agent_dir_respects_pi_config_dir() {
        let dir = omp_agent_dir_with_env(Path::new("/home/test"), |key| {
            (key == "PI_CONFIG_DIR").then(|| OsString::from("custom-omp"))
        });

        assert_eq!(dir, PathBuf::from("/home/test/custom-omp/agent"));
    }

    #[test]
    fn test_omp_agent_dir_respects_pi_coding_agent_dir() {
        let dir = omp_agent_dir_with_env(Path::new("/home/test"), |key| {
            (key == "PI_CODING_AGENT_DIR").then(|| OsString::from("/tmp/omp-agent"))
        });

        assert_eq!(dir, PathBuf::from("/tmp/omp-agent"));
    }

    #[test]
    fn test_uninstall_no_omp_extension_file() {
        let tmp = tempfile::tempdir().unwrap();
        let ext_path = tmp.path().join("extensions/workmux-status.ts");
        let result = uninstall_at(ext_path).unwrap();
        assert!(result.contains("No omp extension found"));
    }

    #[test]
    fn test_uninstall_removes_omp_extension_file() {
        let tmp = tempfile::tempdir().unwrap();
        let ext_dir = tmp.path().join("extensions");
        std::fs::create_dir_all(&ext_dir).unwrap();
        let ext_path = ext_dir.join("workmux-status.ts");
        std::fs::write(&ext_path, "// extension").unwrap();

        let result = uninstall_at(ext_path.clone()).unwrap();
        assert!(result.contains("Removed omp extension"));
        assert!(!ext_path.exists());
    }

    #[test]
    fn test_uninstall_omp_idempotent() {
        let tmp = tempfile::tempdir().unwrap();
        let ext_path = tmp.path().join("extensions/workmux-status.ts");
        let result1 = uninstall_at(ext_path.clone()).unwrap();
        assert!(result1.contains("No omp extension found"));
        let result2 = uninstall_at(ext_path).unwrap();
        assert!(result2.contains("No omp extension found"));
    }
}