flow-mumu 0.1.4

use mumu::parser::interpreter::apply::apply_n_ary_function_value;
use std::sync::{Arc, Mutex};
use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::types::{Value, FunctionValue, IteratorKind};
use crate::common::is_placeholder;

pub fn register_flow_reduce(interp: &mut Interpreter) {
    let f = Arc::new(Mutex::new(flow_reduce_bridge));
    let info = DynamicFnInfo::new(f, true);
    interp.register_dynamic_function_ex("flow:reduce", info);

    interp.set_variable(
        "flow:reduce",
        Value::Function(Box::new(FunctionValue::Named("flow:reduce".to_string())))
    );
}

fn flow_reduce_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_partial(None, None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_partial(None, None, None))
            } else if let Value::Function(fb) = a {
                Ok(make_partial(Some(fb.clone()), None, None))
            } else {
                Err(format!("flow:reduce => first arg must be function or '_', got {:?}", a))
            }
        }
        2 => {
            let a = &args[0];
            let b = &args[1];
            let f_opt = if is_placeholder(a) { None } else if let Value::Function(fb) = a { Some(fb.clone()) } else { return Err(format!("flow:reduce => first arg must be function or '_', got {:?}", a)) };
            let init_opt = if is_placeholder(b) { None } else { Some(b.clone()) };
            Ok(make_partial(f_opt, init_opt, None))
        }
        3 => {
            let f = &args[0];
            let init = &args[1];
            let src = &args[2];
            if is_placeholder(f) || is_placeholder(init) || is_placeholder(src) {
                let f_opt = if let Value::Function(fb) = f { Some(fb.clone()) } else { None };
                let i_opt = if is_placeholder(init) { None } else { Some(init.clone()) };
                let s_opt = if is_placeholder(src) { None } else { Some(src.clone()) };
                Ok(make_partial(f_opt, i_opt, s_opt))
            } else {
                let fb = match f { Value::Function(x) => x.clone(), other => return Err(format!("flow:reduce => first arg must be function, got {:?}", other)) };
                reduce_finalize(interp, fb, init.clone(), src.clone())
            }
        }
        n => Err(format!("flow:reduce => expected up to 3 arguments, got {}", n)),
    }
}

fn make_partial(fn_opt: Option<Box<FunctionValue>>, init_opt: Option<Value>, src_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    Value::Function(Box::new(RustClosure(
        "flow:reduce-partial".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut f_cur = fn_opt.clone();
            let mut i_cur = init_opt.clone();
            let mut s_cur = src_opt.clone();

            for a in new_args {
                if f_cur.is_none() {
                    if is_placeholder(&a) {
                        // remain None
                    } else if let Value::Function(fb) = a {
                        f_cur = Some(fb);
                    } else {
                        return Err("flow:reduce => first arg must be function".to_string());
                    }
                    continue;
                }
                if i_cur.is_none() {
                    if is_placeholder(&a) {
                        // remain None
                    } else {
                        i_cur = Some(a);
                    }
                    continue;
                }
                if s_cur.is_none() {
                    if is_placeholder(&a) {
                        // remain None
                    } else {
                        s_cur = Some(a);
                    }
                    continue;
                }
                return Err("flow:reduce => partial => too many arguments".to_string());
            }

            if f_cur.is_some() && i_cur.is_some() && s_cur.is_some() {
                reduce_finalize(
                    interp,
                    f_cur.as_ref().unwrap().clone(),
                    i_cur.as_ref().unwrap().clone(),
                    s_cur.as_ref().unwrap().clone(),
                )
            } else {
                Ok(make_partial(f_cur, i_cur, s_cur))
            }
        })),
        0,
    )))
}

fn reduce_finalize(interp: &mut Interpreter, reducer: Box<FunctionValue>, init: Value, source: Value) -> Result<Value, String> {
    match source {
        // Iterator source
        Value::Iterator(handle) => {
            let mut acc = init;
            match &handle.kind {
                IteratorKind::Core(state_arc) => {
                    loop {
                        let item_opt = {
                            let mut g = state_arc.lock().map_err(|_| "flow:reduce => iterator lock".to_string())?;
                            if g.done || g.current >= g.end {
                                g.done = true;
                                None
                            } else {
                                let v = Value::Int(g.current);
                                g.current += 1;
                                if g.current >= g.end {
                                    g.done = true;
                                }
                                Some(v)
                            }
                        };
                        match item_opt {
                            Some(v) => {
                                acc = apply_n_ary_function_value(interp, reducer.clone(), vec![acc, v])?;
                            }
                            None => break,
                        }
                    }
                    Ok(acc)
                }
                IteratorKind::Plugin(plugin_arc) => {
                    let mut plugin = plugin_arc.lock().map_err(|_| "flow:reduce => plugin lock".to_string())?;
                    loop {
                        match plugin.next_value() {
                            Ok(v) => {
                                acc = apply_n_ary_function_value(interp, reducer.clone(), vec![acc, v])?;
                            }
                            Err(e) if e == "NO_MORE_DATA" => break,
                            Err(e) => return Err(e),
                        }
                    }
                    Ok(acc)
                }
            }
        }
        // Transform function source — call until NO_MORE_DATA
        Value::Function(tf) => {
            let mut acc = init;
            loop {
                match apply_n_ary_function_value(interp, tf.clone(), vec![]) {
                    Ok(v) => {
                        acc = apply_n_ary_function_value(interp, reducer.clone(), vec![acc, v])?;
                    }
                    Err(e) if e == "NO_MORE_DATA" => break,
                    Err(e) => return Err(e),
                }
            }
            Ok(acc)
        }
        other => Err(format!("flow:reduce => third arg must be Iterator or transform Function, got {:?}", other)),
    }
}