mumusys/

lib.rs

// sys/src/lib.rs
//
// sys plugin for MuMu/Lava:
//   - sys:command(cmdString, callback)
//   - sys:timestamp_ms()
//   - sys:timestamp_micro()
//   - sys:command_stream(cmdString)  [NEW: streaming output as InkIterator]

use mumu::{
    parser::interpreter::{Interpreter, apply_one_function_value},
    parser::types::{FunctionValue, Value, InkIteratorHandle, InkIteratorKind, LavaStream},
};
use std::ffi::c_void;

use std::{
    collections::HashMap,
    process::{Command as SystemCommand, Stdio},
    sync::{
        atomic::{AtomicUsize, Ordering},
        mpsc::{channel, Receiver, Sender, RecvTimeoutError},
        Arc, Mutex,
    },
    thread,
    time::{SystemTime, UNIX_EPOCH, Duration},
    io::{BufRead, BufReader},
};
use lazy_static::lazy_static;
use indexmap::IndexMap;

lazy_static! {
    static ref COMMAND_MANAGER: Mutex<CommandManager> = Mutex::new(CommandManager::new());
}

#[derive(Debug)]
enum CommandMessage {
    Ok(usize, Value),
    Err(usize, String),
}

struct CommandTask {
    callback: Box<FunctionValue>,
    done: bool,
}

struct CommandManager {
    next_id: usize,
    tasks: HashMap<usize, CommandTask>,
    tx: Sender<CommandMessage>,
    rx: Receiver<CommandMessage>,
    active_count: AtomicUsize,
}

impl CommandManager {
    fn new() -> Self {
        let (tx, rx) = channel();
        Self {
            next_id: 0,
            tasks: HashMap::new(),
            tx,
            rx,
            active_count: AtomicUsize::new(0),
        }
    }

    fn add_command_task(&mut self, cmd_str: String, callback: Box<FunctionValue>) {
        let task_id = self.next_id;
        self.next_id += 1;

        self.tasks.insert(
            task_id,
            CommandTask {
                callback,
                done: false,
            },
        );
        self.active_count.fetch_add(1, Ordering::SeqCst);

        let txc = self.tx.clone();
        thread::spawn(move || {
            let parts: Vec<&str> = cmd_str.split_whitespace().collect();
            if parts.is_empty() {
                let _ = txc.send(CommandMessage::Err(task_id, "Empty command".into()));
                return;
            }
            let prog = parts[0];
            let prog_args = &parts[1..];

            match SystemCommand::new(prog).args(prog_args).output() {
                Ok(out) => {
                    let code_i32 = out.status.code().unwrap_or(-1);
                    let success_bool = code_i32 == 0;

                    let stdout_str = String::from_utf8_lossy(&out.stdout).into_owned();
                    let stderr_str = String::from_utf8_lossy(&out.stderr).into_owned();

                    let mut map = IndexMap::new();
                    map.insert("command".into(), Value::SingleString(cmd_str));
                    map.insert("stdout".into(), Value::SingleString(stdout_str));
                    map.insert("stderr".into(), Value::SingleString(stderr_str));
                    map.insert("exit".into(), Value::Int(code_i32));
                    map.insert("success".into(), Value::Bool(success_bool));

                    let val = Value::KeyedArray(map);
                    let _ = txc.send(CommandMessage::Ok(task_id, val));
                }
                Err(e) => {
                    let _ = txc.send(CommandMessage::Err(task_id, format!("Command error: {e}")));
                }
            }
        });
    }

    fn poll_events(&mut self, interp: &mut Interpreter) {
        while let Ok(msg) = self.rx.try_recv() {
            match msg {
                CommandMessage::Ok(id, data_val) => {
                    if let Some(t) = self.tasks.get_mut(&id) {
                        if !t.done {
                            t.done = true;
                            let _ = apply_one_function_value(interp, t.callback.clone(), data_val);
                        }
                    }
                }
                CommandMessage::Err(id, err_str) => {
                    if let Some(t) = self.tasks.get_mut(&id) {
                        if !t.done {
                            t.done = true;
                            let mut map = IndexMap::new();
                            map.insert("command".into(), Value::SingleString("<unknown>".into()));
                            map.insert("stdout".into(), Value::SingleString(String::new()));
                            map.insert("stderr".into(), Value::SingleString(err_str));
                            map.insert("exit".into(), Value::Int(1));
                            map.insert("success".into(), Value::Bool(false));
                            let val = Value::KeyedArray(map);

                            let _ = apply_one_function_value(interp, t.callback.clone(), val);
                        }
                    }
                }
            }
        }

        let before = self.tasks.len();
        self.tasks.retain(|_, task| !task.done);
        let removed = before.saturating_sub(self.tasks.len());
        if removed > 0 {
            self.active_count.fetch_sub(removed, Ordering::SeqCst);
        }
    }

    fn count_active_tasks(&self) -> usize {
        self.active_count.load(Ordering::SeqCst)
    }
}

/* ─────────────────────────── bridge functions ─────────────────────────── */

fn sys_command_bridge(
    _interp: &mut Interpreter,
    mut args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 2 {
        return Err(format!(
            "sys:command expects (cmdString, callback); got {} arg(s)",
            args.len()
        ));
    }

    let cmd_str = match args.remove(0) {
        Value::SingleString(s) => s,
        Value::StrArray(ss) if ss.len() == 1 => ss[0].clone(),
        _ => return Err("sys:command first arg must be a single string".into()),
    };

    let cb_func = match args.remove(0) {
        Value::Function(fb) => fb,
        _ => return Err("sys:command second arg must be a function".into()),
    };

    COMMAND_MANAGER.lock().unwrap().add_command_task(cmd_str, cb_func);

    Ok(Value::Bool(true))
}

fn sys_timestamp_ms_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if !args.is_empty() {
        return Err(format!(
            "sys:timestamp_ms expects 0 arguments, got {}",
            args.len()
        ));
    }
    let now = SystemTime::now();
    let ms = now
        .duration_since(UNIX_EPOCH)
        .map_err(|e| format!("SystemTime error: {e}"))?
        .as_millis();
    Ok(Value::Long(ms as i64))
}

fn sys_timestamp_micro_bridge(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    if !args.is_empty() {
        return Err(format!(
            "sys:timestamp_micro expects 0 arguments, got {}",
            args.len()
        ));
    }
    let now = SystemTime::now();
    let us = now
        .duration_since(UNIX_EPOCH)
        .map_err(|e| format!("SystemTime error: {e}"))?
        .as_micros();
    Ok(Value::Long(us as i64))
}

// -------------- NEW: sys:command_stream(cmdString) --------------

fn sys_command_stream_bridge(
    _interp: &mut Interpreter,
    mut args: Vec<Value>,
) -> Result<Value, String> {
    if args.is_empty() {
        return Err("sys:command_stream expects at least 1 argument".into());
    }

    let cmd_str = match args.remove(0) {
        Value::SingleString(s) => s,
        Value::StrArray(ss) if ss.len() == 1 => ss[0].clone(),
        _ => return Err("sys:command_stream first arg must be a single string".into()),
    };

    // Use system shell for shell features, or split by whitespace for bare commands
    // Here, for full generality, we run: sh -c "cmd_str"
    let (tx, rx): (Sender<String>, Receiver<String>) = channel();

    thread::spawn(move || {
        let mut child = match SystemCommand::new("sh")
            .arg("-c")
            .arg(&cmd_str)
            .stdout(Stdio::piped())
            .spawn()
        {
            Ok(c) => c,
            Err(e) => {
                let _ = tx.send(format!("ERROR: failed to start: {e}"));
                return;
            }
        };

        if let Some(stdout) = child.stdout.take() {
            let reader = BufReader::new(stdout);
            for line in reader.lines() {
                match line {
                    Ok(l) => {
                        if tx.send(l).is_err() {
                            break;
                        }
                    }
                    Err(e) => {
                        let _ = tx.send(format!("ERROR: {e}"));
                        break;
                    }
                }
            }
        }
        let _ = child.wait();
    });

    let state = Arc::new(Mutex::new(CommandStreamState { rx, done: false }));

    let handle = InkIteratorHandle {
        kind: InkIteratorKind::Plugin(Arc::new(Mutex::new(Box::new(CommandStreamPlugin { state: state.clone() })))),
    };

    Ok(Value::InkIterator(handle))
}

struct CommandStreamState {
    rx: Receiver<String>,
    done: bool,
}

struct CommandStreamPlugin {
    state: Arc<Mutex<CommandStreamState>>,
}

impl LavaStream for CommandStreamPlugin {
    fn next_value(&mut self) -> Result<Value, String> {
        let mut state = self.state.lock().unwrap();
        if state.done {
            return Err("NO_MORE_DATA".to_string());
        }
        match state.rx.recv_timeout(Duration::from_millis(50)) {
            Ok(line) => Ok(Value::SingleString(line)),
            Err(RecvTimeoutError::Timeout) => Err("NO_MORE_DATA".to_string()),
            Err(_) => {
                state.done = true;
                Err("NO_MORE_DATA".to_string())
            }
        }
    }
}

/* ───────────────────────── plugin entry point ───────────────────────── */

#[export_name = "Cargo_lock"]
pub unsafe extern "C" fn cargo_lock(
    interp_ptr: *mut c_void,
    _extra_str: *const c_void,
) -> i32 {
    if interp_ptr.is_null() {
        return 1;
    }
    let interp = &mut *(interp_ptr as *mut Interpreter);

    // sys:command
    let cmd_fn = Arc::new(Mutex::new(sys_command_bridge));
    interp.register_dynamic_function("sys:command", cmd_fn);
    interp.set_variable(
        "sys:command",
        Value::Function(Box::new(FunctionValue::Named("sys:command".into()))),
    );

    // sys:timestamp_ms
    let ms_fn = Arc::new(Mutex::new(sys_timestamp_ms_bridge));
    interp.register_dynamic_function("sys:timestamp_ms", ms_fn);
    interp.set_variable(
        "sys:timestamp_ms",
        Value::Function(Box::new(FunctionValue::Named("sys:timestamp_ms".into()))),
    );

    // sys:timestamp_micro
    let micro_fn = Arc::new(Mutex::new(sys_timestamp_micro_bridge));
    interp.register_dynamic_function("sys:timestamp_micro", micro_fn);
    interp.set_variable(
        "sys:timestamp_micro",
        Value::Function(Box::new(FunctionValue::Named("sys:timestamp_micro".into()))),
    );

    // sys:command_stream
    let stream_fn = Arc::new(Mutex::new(sys_command_stream_bridge));
    interp.register_dynamic_function("sys:command_stream", stream_fn);
    interp.set_variable(
        "sys:command_stream",
        Value::Function(Box::new(FunctionValue::Named("sys:command_stream".into()))),
    );

    // background poller (for sys:command)
    let poller = Arc::new(Mutex::new(|interp: &mut Interpreter| {
        let mut mgr = COMMAND_MANAGER.lock().unwrap();
        mgr.poll_events(interp);
        mgr.count_active_tasks()
    }));
    interp.add_poller(poller);

    0
}