array-mumu 0.2.0-rc.5

// array/src/simple.rs
//
// Implements a few array bridges, including a robust `array:filter` that:
//  - Works with IntArray / FloatArray / StrArray / BoolArray / MixedArray
//  - Supports *partial application* and "_" placeholder completion
//  - Returns **named dynamic functions** for partials (avoids FFI RustClosure crashes)
//
// Other helpers included here (kept for compatibility with existing registrations):
//  - array:zip (with partial support via RustClosure)
//  - cmp_value (used by sort helpers)
//  - array:sort (local variant; main export may come from src/sort.rs in this crate)
//  - array:flatten
//  - array:join
//  - array:zip_with
//  - array:prop
//
// Note: This file uses only public interpreter API via `apply_*` helpers to
// call functions supplied by the user from Lava/MuMu code.

use mumu::{
    parser::interpreter::Interpreter,
    parser::types::{FunctionValue, Value},
};
use crate::apply::{apply_one_function_value, apply_n_ary_function_value};

use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicU64, Ordering};

static GENSYM: AtomicU64 = AtomicU64::new(1);
fn gensym(prefix: &str) -> String {
    let n = GENSYM.fetch_add(1, Ordering::Relaxed);
    format!("__array_filter_partial_{}_{}", prefix, n)
}

fn is_placeholder(v: &Value) -> bool {
    matches!(v, Value::Placeholder)
        || matches!(v, Value::SingleString(s) if s == "_")
        || matches!(v, Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_")
}

fn looks_like_function(v: &Value) -> bool {
    matches!(v, Value::Function(_))
}

/* ──────────────────────────────────────────────────────────────────────────
   array:filter – MixedArray/object support + **named partials**
   ──────────────────────────────────────────────────────────────────────── */

pub fn array_filter(
    interp: &mut Interpreter,
    mut args: Vec<Value>,
) -> Result<Value, String> {
    match args.len() {
        // Fully partial: wait for (fn, data)
        0 => make_named_partial_filter(interp, None, None),

        // One argument: could be `_`, a function, or data
        1 => {
            let a = args.remove(0);
            if is_placeholder(&a) {
                make_named_partial_filter(interp, None, None)
            } else if looks_like_function(&a) {
                // fn known, data unknown
                make_named_partial_filter(interp, Some(a), None)
            } else {
                // treat as data known, wait for fn
                make_named_partial_filter(interp, None, Some(a))
            }
        }

        // Two arguments: placeholders allowed; or run immediately if both real
        2 => {
            let f = args.remove(0);
            let d = args.remove(0);

            let f_pl = is_placeholder(&f);
            let d_pl = is_placeholder(&d);

            if f_pl && d_pl {
                return make_named_partial_filter(interp, None, None);
            }
            if f_pl {
                return make_named_partial_filter(interp, None, Some(d));
            }
            if d_pl {
                return make_named_partial_filter(interp, Some(f), None);
            }

            // both real => run now
            do_filter(interp, f, d)
        }

        n => Err(format!("array:filter expects up to 2 arguments: function, array (got {})", n)),
    }
}

/// Register a *named* dynamic function that completes the filter call.
/// Returning Named avoids cross-FFI RustClosure lifetime issues.
fn make_named_partial_filter(
    interp: &mut Interpreter,
    f_opt: Option<Value>,
    d_opt: Option<Value>,
) -> Result<Value, String> {
    let name = gensym("f");

    // capture current partial state
    let c_f = f_opt.clone();
    let c_d = d_opt.clone();

    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| -> Result<Value, String> {
        let mut f = c_f.clone();
        let mut d = c_d.clone();

        for arg in new_args {
            if f.is_none() && looks_like_function(&arg) && !is_placeholder(&arg) {
                f = Some(arg);
                continue;
            }
            if d.is_none() && !is_placeholder(&arg) {
                d = Some(arg);
                continue;
            }
            return Err("array:filter => partial => too many or invalid arguments".to_string());
        }

        if let (Some(ff), Some(dd)) = (f.clone(), d.clone()) {
            return do_filter(interp, ff, dd);
        }

        // still partial => chain another named partial
        make_named_partial_filter(interp, f, d)
    };

    let dyn_fn = Arc::new(Mutex::new(closure));
    interp.register_dynamic_function(&name, dyn_fn);

    Ok(Value::Function(Box::new(FunctionValue::Named(name))))
}

fn do_filter(
    interp: &mut Interpreter,
    func: Value,
    arr: Value,
) -> Result<Value, String> {
    match arr {
        Value::IntArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::Int(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate must return bool".to_string()),
                }
            }
            Ok(Value::IntArray(out))
        }
        Value::FloatArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::Float(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate must return bool".to_string()),
                }
            }
            Ok(Value::FloatArray(out))
        }
        Value::StrArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::SingleString(x.clone()))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate must return bool".to_string()),
                }
            }
            Ok(Value::StrArray(out))
        }
        Value::BoolArray(xs) => {
            let mut out = Vec::new();
            for x in xs {
                let res = match &func {
                    Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::Bool(x))?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(x),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate must return bool".to_string()),
                }
            }
            Ok(Value::BoolArray(out))
        }
        // MixedArray/object arrays (e.g., KeyedArray items)
        Value::MixedArray(xs) => {
            let mut out = Vec::new();
            for v in xs {
                let res = match &func {
                    Value::Function(fb) => apply_one_function_value(interp, fb.clone(), v.clone())?,
                    _ => return Err("array:filter: first arg must be function".to_string()),
                };
                match res {
                    Value::Bool(true) => out.push(v),
                    Value::Bool(false) => {}
                    _ => return Err("array:filter: predicate must return bool".to_string()),
                }
            }
            Ok(Value::MixedArray(out))
        }
        _ => Err("array:filter: only IntArray, FloatArray, StrArray, BoolArray, or MixedArray supported".to_string()),
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   array:zip (with partial and placeholder support)
   ──────────────────────────────────────────────────────────────────────── */

fn is_placeholder_for_zip(v: &Value) -> bool {
    is_placeholder(v)
}

pub fn array_zip(
    _interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_zip_partial(None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder_for_zip(a) {
                Ok(make_zip_partial(None, None))
            } else {
                Ok(make_zip_partial(Some(a.clone()), None))
            }
        }
        2 => {
            let a1 = &args[0];
            let a2 = &args[1];
            let a1_is_pl = is_placeholder_for_zip(a1);
            let a2_is_pl = is_placeholder_for_zip(a2);
            if a1_is_pl || a2_is_pl {
                Ok(make_zip_partial(
                    if a1_is_pl { None } else { Some(a1.clone()) },
                    if a2_is_pl { None } else { Some(a2.clone()) },
                ))
            } else {
                do_zip(a1.clone(), a2.clone())
            }
        }
        n => Err(format!("array:zip expects up to 2 arguments, got {}", n)),
    }
}

fn do_zip(a: Value, b: Value) -> Result<Value, String> {
    match (a, b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                out.push(vec![xs[i], ys[i]]);
            }
            Ok(Value::Int2DArray(out))
        }
        (Value::StrArray(xs), Value::StrArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                out.push(vec![xs[i].clone(), ys[i].clone()]);
            }
            // There is NO Str2DArray; wrap each as StrArray for MixedArray
            Ok(Value::MixedArray(out.into_iter().map(Value::StrArray).collect()))
        }
        _ => Err("array:zip: only IntArray/IntArray or StrArray/StrArray".to_string()),
    }
}

fn make_zip_partial(a_opt: Option<Value>, b_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    let closure = move |_interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut a = a_opt.clone();
        let mut b = b_opt.clone();

        for arg in new_args {
            if a.is_none() {
                if is_placeholder(&arg) {
                    // remain None
                } else {
                    a = Some(arg);
                }
                continue;
            }
            if b.is_none() {
                if is_placeholder(&arg) {
                    // remain None
                } else {
                    b = Some(arg);
                }
                continue;
            }
            return Err("array:zip partial: too many arguments".to_string());
        }

        if a.is_some() && b.is_some() {
            do_zip(a.unwrap(), b.unwrap())
        } else {
            Ok(make_zip_partial(a, b))
        }
    };

    Value::Function(Box::new(RustClosure(
        "array:zip-partial".to_string(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

/* ──────────────────────────────────────────────────────────────────────────
   Misc helpers retained for compatibility / usage by other bridges
   ──────────────────────────────────────────────────────────────────────── */

pub fn cmp_value(a: &Value, b: &Value) -> std::cmp::Ordering {
    match (a, b) {
        (Value::Int(a), Value::Int(b)) => a.cmp(b),
        (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal),
        (Value::SingleString(a), Value::SingleString(b)) => a.cmp(b),
        (Value::StrArray(a), Value::StrArray(b)) => a.cmp(b),
        _ => std::cmp::Ordering::Equal,
    }
}

pub fn array_sort(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 1 && args.len() != 2 {
        return Err("array:sort expects array and optional function".to_string());
    }
    let arr = args.remove(0);
    if args.is_empty() {
        match arr {
            Value::IntArray(mut xs) => {
                xs.sort();
                Ok(Value::IntArray(xs))
            }
            Value::StrArray(mut xs) => {
                xs.sort();
                Ok(Value::StrArray(xs))
            }
            Value::FloatArray(mut xs) => {
                xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
                Ok(Value::FloatArray(xs))
            }
            Value::MixedArray(items) => {
                if items.iter().all(|v| matches!(v, Value::Int(_))) {
                    let mut vals: Vec<i32> = items.iter().map(|v| match v { Value::Int(i) => *i, _ => 0 }).collect();
                    vals.sort();
                    Ok(Value::IntArray(vals))
                } else if items.iter().all(|v| matches!(v, Value::Float(_))) {
                    let mut vals: Vec<f64> = items.iter().map(|v| match v { Value::Float(f) => *f, _ => 0.0 }).collect();
                    vals.sort_by(|a, b| a.partial_cmp(b).unwrap());
                    Ok(Value::FloatArray(vals))
                } else if items.iter().all(|v| matches!(v, Value::SingleString(_))) {
                    let mut vals: Vec<String> = items.iter().map(|v| match v { Value::SingleString(s) => s.clone(), _ => "".to_string() }).collect();
                    vals.sort();
                    Ok(Value::StrArray(vals))
                } else {
                    Ok(Value::MixedArray(items))
                }
            }
            _ => Err("array:sort: only IntArray, FloatArray, StrArray, or MixedArray".to_string()),
        }
    } else {
        let func = args.remove(0);
        match arr {
            Value::IntArray(xs) => {
                let mut decorated: Vec<(i32, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::Int(*x))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((*x, key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::IntArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::StrArray(xs) => {
                let mut decorated: Vec<(String, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::SingleString(x.clone()))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((x.clone(), key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::StrArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::FloatArray(xs) => {
                let mut decorated: Vec<(f64, Value)> = Vec::new();
                for x in &xs {
                    let key = match &func {
                        Value::Function(fb) => apply_one_function_value(interp, fb.clone(), Value::Float(*x))?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((*x, key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                Ok(Value::FloatArray(decorated.into_iter().map(|(x,_)| x).collect()))
            }
            Value::MixedArray(items) => {
                // Use key function, but preserve result type if possible
                let mut decorated: Vec<(Value, Value)> = Vec::new();
                for item in &items {
                    let key = match &func {
                        Value::Function(fb) => apply_one_function_value(interp, fb.clone(), item.clone())?,
                        _ => return Err("array:sort: 2nd arg must be function".to_string()),
                    };
                    decorated.push((item.clone(), key));
                }
                decorated.sort_by(|a, b| cmp_value(&a.1, &b.1));
                let values: Vec<Value> = decorated.into_iter().map(|(v, _)| v).collect();
                // Coerce if all elements of the same type
                if values.iter().all(|v| matches!(v, Value::Int(_))) {
                    Ok(Value::IntArray(values.into_iter().map(|v| match v { Value::Int(i) => i, _ => 0 }).collect()))
                } else if values.iter().all(|v| matches!(v, Value::Float(_))) {
                    Ok(Value::FloatArray(values.into_iter().map(|v| match v { Value::Float(f) => f, _ => 0.0 }).collect()))
                } else if values.iter().all(|v| matches!(v, Value::SingleString(_))) {
                    Ok(Value::StrArray(values.into_iter().map(|v| match v { Value::SingleString(s) => s, _ => "".to_string() }).collect()))
                } else {
                    Ok(Value::MixedArray(values))
                }
            }
            _ => Err("array:sort: only IntArray, FloatArray, StrArray, MixedArray".to_string()),
        }
    }
}

// SHALLOW flatten: flattens one level only.
pub fn array_flatten(_interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 1 {
        return Err("array:flatten expects one array".to_string());
    }
    let arr = args.remove(0);
    match arr {
        Value::IntArray(xs) => Ok(Value::IntArray(xs)),
        Value::FloatArray(xs) => Ok(Value::FloatArray(xs)),
        Value::StrArray(xs) => Ok(Value::StrArray(xs)),
        Value::BoolArray(xs) => Ok(Value::BoolArray(xs)),
        Value::MixedArray(items) => {
            let mut flat = Vec::new();
            for v in items {
                match v {
                    Value::IntArray(x) => for i in x { flat.push(Value::Int(i)); }
                    Value::FloatArray(f) => for ff in f { flat.push(Value::Float(ff)); }
                    Value::StrArray(s) => for ss in s { flat.push(Value::SingleString(ss)); }
                    Value::BoolArray(b) => for bb in b { flat.push(Value::Bool(bb)); }
                    Value::MixedArray(nested) => flat.push(Value::MixedArray(nested)),
                    other => flat.push(other),
                }
            }
            Ok(Value::MixedArray(flat))
        }
        _ => Err("array:flatten: unsupported type".to_string()),
    }
}

pub fn array_join(_interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 2 {
        return Err("array:join expects 2 arguments".to_string());
    }
    let arr = args.remove(0);
    let delim = args.remove(0);
    let delim = match delim {
        Value::SingleString(s) => s,
        Value::StrArray(xs) if xs.len() == 1 => xs[0].clone(),
        _ => return Err("array:join => second argument must be StrArray or string".to_string()),
    };
    match arr {
        Value::StrArray(xs) => Ok(Value::SingleString(xs.join(&delim))),
        Value::IntArray(xs) => Ok(Value::SingleString(xs.iter().map(|i| i.to_string()).collect::<Vec<_>>().join(&delim))),
        Value::FloatArray(xs) => Ok(Value::SingleString(xs.iter().map(|f| f.to_string()).collect::<Vec<_>>().join(&delim))),
        Value::BoolArray(xs) => Ok(Value::SingleString(xs.iter().map(|b| b.to_string()).collect::<Vec<_>>().join(&delim))),
        _ => Err("array:join => first argument must be StrArray, IntArray, FloatArray, or BoolArray".to_string()),
    }
}

pub fn array_zipwith(interp: &mut Interpreter, mut args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 3 {
        return Err("array:zip_with expects function and two arrays".to_string());
    }
    let func = args.remove(0);
    let a = args.remove(0);
    let b = args.remove(0);
    match (a, b) {
        (Value::IntArray(xs), Value::IntArray(ys)) => {
            let n = xs.len().min(ys.len());
            let mut out = Vec::with_capacity(n);
            for i in 0..n {
                let res = match &func {
                    Value::Function(fb) => apply_n_ary_function_value(
                        interp, fb.clone(), vec![Value::Int(xs[i]), Value::Int(ys[i])]
                    )?,
                    _ => return Err("array:zip_with: first arg must be function".to_string()),
                };
                match res {
                    Value::Int(z) => out.push(z),
                    _ => return Err("array:zip_with: function must return int".to_string()),
                }
            }
            Ok(Value::IntArray(out))
        }
        _ => Err("array:zip_with: only IntArray/IntArray supported".to_string()),
    }
}

pub fn array_prop(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    if args.len() != 2 {
        return Err("array:prop expects 2 args".to_string());
    }
    let key = match &args[0] {
        Value::SingleString(s) => s,
        Value::StrArray(ss) if ss.len() == 1 => &ss[0],
        _ => return Err("array:prop expects a string key as the first argument".to_string()),
    };
    let arr = &args[1];
    match arr {
        Value::KeyedArray(map) => map.get(key).cloned().ok_or_else(|| format!("Key '{}' not found", key)),
        _ => Err("array:prop expects keyed array as the second argument".to_string()),
    }
}