array-mumu 0.2.0-rc.5

// src/sort.rs
//
// Robust sorting bridges:
//   - array:sort([cmp?], array?)            → default ascending or custom comparator
//   - array:sort_by(keyFn, array?)         → key extractor, stable
//   - array:sort_with([cmp...], array?)    → chain of comparators, stable
//
// Features:
//   • Full partial-application and "_" placeholder support
//   • Typed outputs: IntArray / FloatArray / StrArray when homogeneous
//   • Accept homogeneous MixedArray for convenience (coerces to typed on return)
//   • Stable sorts so equal comps preserve input order
//
// Comparator conventions:
//   - Custom comparator must return negative / zero / positive (Int/Long/Float).
//   - Non-numeric or Placeholder results are treated as 0 (tie).

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

use std::cmp::Ordering;
use std::sync::{Arc, Mutex};

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] == "_")
}

/* ──────────────────────────────────────────────────────────────────────────
   Utilities: gather / return typed arrays
   ─────────────────────────────────────────────────────────────────────── */

#[derive(Clone)]
enum Homog {
    Int(Vec<i32>),
    Float(Vec<f64>),
    Str(Vec<String>),
    Mixed(Vec<Value>),
}

fn to_homog(v: &Value) -> Result<Homog, String> {
    match v {
        Value::IntArray(xs)   => Ok(Homog::Int(xs.clone())),
        Value::FloatArray(xs) => Ok(Homog::Float(xs.clone())),
        Value::StrArray(xs)   => Ok(Homog::Str(xs.clone())),
        Value::MixedArray(xs) => {
            if xs.iter().all(|x| matches!(x, Value::Int(_))) {
                Ok(Homog::Int(xs.iter().map(|x| if let Value::Int(i)=x { *i } else { unreachable!() }).collect()))
            } else if xs.iter().all(|x| matches!(x, Value::Float(_))) {
                Ok(Homog::Float(xs.iter().map(|x| if let Value::Float(f)=x { *f } else { unreachable!() }).collect()))
            } else if xs.iter().all(|x| matches!(x, Value::SingleString(_))) {
                Ok(Homog::Str(xs.iter().map(|x| if let Value::SingleString(s)=x { s.clone() } else { unreachable!() }).collect()))
            } else {
                Ok(Homog::Mixed(xs.clone()))
            }
        }
        other => Err(format!("array:sort => unsupported array type: {:?}", other)),
    }
}

fn from_homog(h: Homog) -> Value {
    match h {
        Homog::Int(v)   => Value::IntArray(v),
        Homog::Float(v) => Value::FloatArray(v),
        Homog::Str(v)   => Value::StrArray(v),
        Homog::Mixed(v) => Value::MixedArray(v),
    }
}

fn typed_from_values(vals: Vec<Value>) -> Value {
    if vals.iter().all(|x| matches!(x, Value::Int(_))) {
        Value::IntArray(vals.into_iter().map(|x| if let Value::Int(i)=x { i } else { unreachable!() }).collect())
    } else if vals.iter().all(|x| matches!(x, Value::Float(_))) {
        Value::FloatArray(vals.into_iter().map(|x| if let Value::Float(f)=x { f } else { unreachable!() }).collect())
    } else if vals.iter().all(|x| matches!(x, Value::SingleString(_))) {
        Value::StrArray(vals.into_iter().map(|x| if let Value::SingleString(s)=x { s } else { unreachable!() }).collect())
    } else {
        Value::MixedArray(vals)
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   Comparator helpers
   ─────────────────────────────────────────────────────────────────────── */

fn ord_from_cmp_val(interp: &mut Interpreter, cmp_fn: &FunctionValue, a: Value, b: Value) -> Ordering {
    // cmp(a, b) -> numeric
    let res = apply_n_ary_function_value(interp, Box::new(cmp_fn.clone()), vec![a, b]);
    let n = match res {
        Ok(Value::Int(i))   => i as f64,
        Ok(Value::Long(l))  => l as f64,
        Ok(Value::Float(f)) => f,
        _                   => 0.0, // treat invalid as tie
    };
    if n < 0.0 { Ordering::Less } else if n > 0.0 { Ordering::Greater } else { Ordering::Equal }
}

/* ──────────────────────────────────────────────────────────────────────────
   array:sort  —  default or custom comparator
   ─────────────────────────────────────────────────────────────────────── */

pub fn array_sort_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial_sort(None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_two_arg_partial_sort(None, None))
            } else if matches!(a, Value::Function(_)) {
                Ok(make_two_arg_partial_sort(Some(a.clone()), None))
            } else {
                // one-arg form: array only → default sort
                do_sort_default_or_cmp(interp, None, a.clone())
            }
        }
        2 => {
            let a = &args[0];
            let b = &args[1];
            let a_pl = is_placeholder(a);
            let b_pl = is_placeholder(b);
            if a_pl || b_pl {
                Ok(make_two_arg_partial_sort(
                    if a_pl { None } else { Some(a.clone()) },
                    if b_pl { None } else { Some(b.clone()) },
                ))
            } else {
                let cmp = if matches!(a, Value::Function(_)) { Some(a.clone()) } else { None };
                let arr = if cmp.is_some() { b.clone() } else { a.clone() };
                do_sort_default_or_cmp(interp, cmp, arr)
            }
        }
        n => Err(format!("array:sort => expected up to 2 args, got {}", n)),
    }
}

fn make_two_arg_partial_sort(cmp_opt: Option<Value>, arr_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut c = cmp_opt.clone();
        let mut d = arr_opt.clone();
        for arg in new_args {
            if c.is_none() {
                if is_placeholder(&arg) { /* stay */ } else if matches!(arg, Value::Function(_)) { c = Some(arg); }
                else { d = Some(arg); }
                continue;
            }
            if d.is_none() {
                if is_placeholder(&arg) { /* stay */ } else { d = Some(arg); }
                continue;
            }
            return Err("array:sort => partial => too many arguments".to_string());
        }
        if let Some(arr) = d.clone() {
            do_sort_default_or_cmp(interp, c.clone(), arr)
        } else {
            Ok(make_two_arg_partial_sort(c, d))
        }
    };
    Value::Function(Box::new(RustClosure(
        "array:sort-partial".into(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

fn do_sort_default_or_cmp(interp: &mut Interpreter, cmp_opt: Option<Value>, arr_val: Value) -> Result<Value, String> {
    let data = to_homog(&arr_val)?;
    match (cmp_opt, data) {
        (None, Homog::Int(mut v)) => { v.sort(); Ok(Value::IntArray(v)) }
        (None, Homog::Float(mut v)) => { v.sort_by(|a,b| a.partial_cmp(b).unwrap()); Ok(Value::FloatArray(v)) }
        (None, Homog::Str(mut v)) => { v.sort(); Ok(Value::StrArray(v)) }
        (None, Homog::Mixed(v)) => Ok(Value::MixedArray(v)), // leave heterogeneous as-is

        (Some(Value::Function(fb)), Homog::Int(v0)) => {
            let mut v: Vec<i32> = v0;
            v.sort_by(|a,b| ord_from_cmp_val(interp, &fb, Value::Int(*a), Value::Int(*b)));
            Ok(Value::IntArray(v))
        }
        (Some(Value::Function(fb)), Homog::Float(v0)) => {
            let mut v: Vec<f64> = v0;
            v.sort_by(|a,b| ord_from_cmp_val(interp, &fb, Value::Float(*a), Value::Float(*b)));
            Ok(Value::FloatArray(v))
        }
        (Some(Value::Function(fb)), Homog::Str(v0)) => {
            let mut v: Vec<String> = v0;
            v.sort_by(|a,b| ord_from_cmp_val(interp, &fb, Value::SingleString(a.clone()), Value::SingleString(b.clone())));
            Ok(Value::StrArray(v))
        }
        (Some(Value::Function(fb)), Homog::Mixed(mut v)) => {
            v.sort_by(|a,b| ord_from_cmp_val(interp, &fb, a.clone(), b.clone()));
            Ok(typed_from_values(v))
        }
        (Some(other), _) => Err(format!("array:sort => first argument must be Function, got {:?}", other)),
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   array:sort_by
   ─────────────────────────────────────────────────────────────────────── */

pub fn array_sort_by_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial_sort_by(None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_two_arg_partial_sort_by(None, None))
            } else if matches!(a, Value::Function(_)) {
                Ok(make_two_arg_partial_sort_by(Some(a.clone()), None))
            } else {
                Err("array:sort_by => first argument must be a function".into())
            }
        }
        2 => {
            let f = &args[0];
            let arr = &args[1];
            let f_pl = is_placeholder(f);
            let a_pl = is_placeholder(arr);
            if f_pl || a_pl {
                Ok(make_two_arg_partial_sort_by(
                    if f_pl { None } else { Some(f.clone()) },
                    if a_pl { None } else { Some(arr.clone()) },
                ))
            } else {
                do_sort_by(interp, f.clone(), arr.clone())
            }
        }
        n => Err(format!("array:sort_by => expected 2 args, got {}", n)),
    }
}

fn make_two_arg_partial_sort_by(f_opt: Option<Value>, arr_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut f = f_opt.clone();
        let mut d = arr_opt.clone();
        for arg in new_args {
            if f.is_none() {
                if is_placeholder(&arg) { /* stay */ } else if matches!(arg, Value::Function(_)) { f = Some(arg); }
                else { return Err("array:sort_by => first argument must be a function".into()); }
                continue;
            }
            if d.is_none() {
                if is_placeholder(&arg) { /* stay */ } else { d = Some(arg); }
                continue;
            }
            return Err("array:sort_by => partial => too many arguments".to_string());
        }
        if let (Some(ff), Some(arr)) = (f.clone(), d.clone()) {
            do_sort_by(interp, ff, arr)
        } else {
            Ok(make_two_arg_partial_sort_by(f, d))
        }
    };
    Value::Function(Box::new(RustClosure(
        "array:sort_by-partial".into(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

fn do_sort_by(interp: &mut Interpreter, key_fn_val: Value, arr_val: Value) -> Result<Value, String> {
    let key_fn = match key_fn_val {
        Value::Function(fb) => fb,
        other => return Err(format!("array:sort_by => first argument must be a function, got {:?}", other)),
    };
    let data = to_homog(&arr_val)?;
    match data {
        Homog::Int(v0) => {
            let mut pairs: Vec<(i32, Value)> = v0.into_iter().map(|x| (x, Value::Int(x))).collect();
            pairs.sort_by(|a, b| {
                let ka = apply_n_ary_function_value(interp, key_fn.clone(), vec![a.1.clone()]).ok();
                let kb = apply_n_ary_function_value(interp, key_fn.clone(), vec![b.1.clone()]).ok();
                ord_from_keys(ka, kb)
            });
            Ok(Value::IntArray(pairs.into_iter().map(|(x,_)| x).collect()))
        }
        Homog::Float(v0) => {
            let mut pairs: Vec<(f64, Value)> = v0.into_iter().map(|x| (x, Value::Float(x))).collect();
            pairs.sort_by(|a, b| {
                let ka = apply_n_ary_function_value(interp, key_fn.clone(), vec![a.1.clone()]).ok();
                let kb = apply_n_ary_function_value(interp, key_fn.clone(), vec![b.1.clone()]).ok();
                ord_from_keys(ka, kb)
            });
            Ok(Value::FloatArray(pairs.into_iter().map(|(x,_)| x).collect()))
        }
        Homog::Str(v0) => {
            let mut pairs: Vec<(String, Value)> = v0.into_iter().map(|x| (x.clone(), Value::SingleString(x))).collect();
            pairs.sort_by(|a, b| {
                let ka = apply_n_ary_function_value(interp, key_fn.clone(), vec![a.1.clone()]).ok();
                let kb = apply_n_ary_function_value(interp, key_fn.clone(), vec![b.1.clone()]).ok();
                ord_from_keys(ka, kb)
            });
            Ok(Value::StrArray(pairs.into_iter().map(|(x,_)| x).collect()))
        }
        Homog::Mixed(mut v) => {
            v.sort_by(|a, b| {
                let ka = apply_n_ary_function_value(interp, key_fn.clone(), vec![a.clone()]).ok();
                let kb = apply_n_ary_function_value(interp, key_fn.clone(), vec![b.clone()]).ok();
                ord_from_keys(ka, kb)
            });
            Ok(typed_from_values(v))
        }
    }
}

fn ord_from_keys(ka: Option<Value>, kb: Option<Value>) -> Ordering {
    use Value::*;
    match (ka, kb) {
        (Some(Int(a)),   Some(Int(b)))   => a.cmp(&b),
        (Some(Long(a)),  Some(Long(b)))  => a.cmp(&b),
        (Some(Float(a)), Some(Float(b))) => a.partial_cmp(&b).unwrap_or(Ordering::Equal),
        (Some(SingleString(a)), Some(SingleString(b))) => a.cmp(&b),
        // Mismatched or invalid → tie
        _ => Ordering::Equal,
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   array:sort_with
   ─────────────────────────────────────────────────────────────────────── */

pub fn array_sort_with_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_two_arg_partial_sort_with(Vec::new(), None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_two_arg_partial_sort_with(Vec::new(), None))
            } else {
                let comps = parse_comparators(a.clone())?;
                Ok(make_two_arg_partial_sort_with(comps, None))
            }
        }
        2 => {
            let a = &args[0];
            let b = &args[1];
            let a_pl = is_placeholder(a);
            let b_pl = is_placeholder(b);
            if a_pl || b_pl {
                let comps = if a_pl { Vec::new() } else { parse_comparators(a.clone())? };
                Ok(make_two_arg_partial_sort_with(comps, if b_pl { None } else { Some(b.clone()) }))
            } else {
                let comps = parse_comparators(a.clone())?;
                do_sort_with(interp, comps, b.clone())
            }
        }
        n => Err(format!("array:sort_with => expected up to 2 args, got {}", n)),
    }
}

fn parse_comparators(val: Value) -> Result<Vec<FunctionValue>, String> {
    match val {
        Value::Function(fb) => Ok(vec![*fb]),
        Value::MixedArray(items) => {
            let mut comps = Vec::new();
            for it in items {
                match it {
                    Value::Function(fb) => comps.push(*fb),
                    other => return Err(format!("array:sort_with => comparator list must be functions, got {:?}", other)),
                }
            }
            Ok(comps)
        }
        other => Err(format!("array:sort_with => first argument must be Function or array of Functions, got {:?}", other)),
    }
}

fn make_two_arg_partial_sort_with(comps: Vec<FunctionValue>, arr_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;
    let closure = move |interp: &mut Interpreter, new_args: Vec<Value>| {
        let mut c = comps.clone();
        let mut d = arr_opt.clone();

        for arg in new_args {
            if c.is_empty() {
                match arg.clone() {
                    Value::Function(fb) => c.push(*fb),
                    Value::MixedArray(items) => {
                        for it in items {
                            if let Value::Function(fb) = it {
                                c.push(*fb);
                            } else if !is_placeholder(&it) {
                                return Err("array:sort_with => comparator list must be functions".to_string());
                            }
                        }
                    }
                    _ if is_placeholder(&arg) => { /* keep empty */ }
                    other => { d = Some(other); }
                }
                continue;
            }
            if d.is_none() {
                if is_placeholder(&arg) { /*keep*/ } else { d = Some(arg); }
                continue;
            }
            return Err("array:sort_with => partial => too many arguments".to_string());
        }

        if let Some(arr) = d.clone() {
            do_sort_with(interp, c.clone(), arr)
        } else {
            Ok(make_two_arg_partial_sort_with(c, d))
        }
    };
    Value::Function(Box::new(RustClosure(
        "array:sort_with-partial".into(),
        Arc::new(Mutex::new(closure)),
        0,
    )))
}

fn do_sort_with(interp: &mut Interpreter, comps: Vec<FunctionValue>, arr_val: Value) -> Result<Value, String> {
    if comps.is_empty() {
        return Err("array:sort_with => requires at least one comparator".into());
    }
    let data = to_homog(&arr_val)?;

    match data {
        Homog::Int(v0) => {
            let mut v: Vec<i32> = v0;
            v.sort_by(|a, b| {
                for cf in &comps {
                    let o = ord_from_cmp_val(interp, cf, Value::Int(*a), Value::Int(*b));
                    if o != Ordering::Equal { return o; }
                }
                Ordering::Equal
            });
            Ok(Value::IntArray(v))
        }
        Homog::Float(v0) => {
            let mut v: Vec<f64> = v0;
            v.sort_by(|a, b| {
                for cf in &comps {
                    let o = ord_from_cmp_val(interp, cf, Value::Float(*a), Value::Float(*b));
                    if o != Ordering::Equal { return o; }
                }
                Ordering::Equal
            });
            Ok(Value::FloatArray(v))
        }
        Homog::Str(v0) => {
            let mut v: Vec<String> = v0;
            v.sort_by(|a, b| {
                for cf in &comps {
                    let o = ord_from_cmp_val(
                        interp,
                        cf,
                        Value::SingleString(a.clone()),
                        Value::SingleString(b.clone()),
                    );
                    if o != Ordering::Equal { return o; }
                }
                Ordering::Equal
            });
            Ok(Value::StrArray(v))
        }
        Homog::Mixed(mut v) => {
            v.sort_by(|a, b| {
                for cf in &comps {
                    let o = ord_from_cmp_val(interp, cf, a.clone(), b.clone());
                    if o != Ordering::Equal { return o; }
                }
                Ordering::Equal
            });
            Ok(typed_from_values(v))
        }
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   Helpers for converting arrays for key/comparator sorting
   ─────────────────────────────────────────────────────────────────────── */

fn to_mixed_vec(arr: &Value) -> Result<Vec<Value>, String> {
    match arr {
        Value::MixedArray(xs) => Ok(xs.clone()),
        Value::IntArray(xs) => Ok(xs.iter().map(|&x| Value::Int(x)).collect()),
        Value::FloatArray(xs) => Ok(xs.iter().map(|&x| Value::Float(x)).collect()),
        Value::StrArray(xs) => Ok(xs.iter().map(|x| Value::SingleString(x.clone())).collect()),
        _ => Err("array:sort_with/sort_by => unsupported array type".to_string()),
    }
}