flow-mumu 0.1.3

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};
use crate::common::{is_placeholder, to_transform};

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

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

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

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

// ---------------- ap_zip (zipped) ----------------

fn flow_ap_zip_bridge(_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(a) {
                Ok(make_zip_partial(None, None))
            } else {
                Ok(make_zip_partial(Some(a.clone()), None))
            }
        }
        2 => {
            let fsrc = &args[0];
            let xsrc = &args[1];
            let f_pl = is_placeholder(fsrc);
            let x_pl = is_placeholder(xsrc);
            if !f_pl && !x_pl {
                ap_zip_finalize(fsrc.clone(), xsrc.clone())
            } else {
                Ok(make_zip_partial(
                    if f_pl { None } else { Some(fsrc.clone()) },
                    if x_pl { None } else { Some(xsrc.clone()) },
                ))
            }
        }
        n => Err(format!("flow:ap_zip => expected up to 2 arguments, got {}", n)),
    }
}

fn make_zip_partial(fsrc_opt: Option<Value>, xsrc_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    Value::Function(Box::new(RustClosure(
        "flow:ap_zip-partial".to_string(),
        Arc::new(Mutex::new(move |_interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut fsrc = fsrc_opt.clone();
            let mut xsrc = xsrc_opt.clone();

            for a in new_args {
                if fsrc.is_none() {
                    if !is_placeholder(&a) { fsrc = Some(a); }
                    continue;
                }
                if xsrc.is_none() {
                    if !is_placeholder(&a) { xsrc = Some(a); }
                    continue;
                }
                return Err("flow:ap_zip => partial => too many arguments".to_string());
            }

            if fsrc.is_some() && xsrc.is_some() {
                ap_zip_finalize(fsrc.unwrap(), xsrc.unwrap())
            } else {
                Ok(make_zip_partial(fsrc, xsrc))
            }
        })),
        0,
    )))
}

fn ap_zip_finalize(fsrc: Value, xsrc: Value) -> Result<Value, String> {
    let fs = to_transform(fsrc)?;
    let xs = to_transform(xsrc)?;

    use mumu::parser::types::FunctionValue::RustClosure;

    let st = Arc::new(Mutex::new(ApZipState { fs, xs }));

    let clos = RustClosure(
        "flow:ap_zip".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            // Clone transforms outside of calls
            let (fs_clone, xs_clone) = {
                let st_guard = st.lock().map_err(|_| "flow:ap_zip => state lock".to_string())?;
                (st_guard.fs.clone(), st_guard.xs.clone())
            };

            // Pull one function
            let f_val = apply_n_ary_function_value(interp, fs_clone, vec![]);
            let f_val = match f_val {
                Ok(v) => v,
                Err(e) => return Err(e),
            };
            let f_fn = match f_val {
                Value::Function(fb) => fb,
                other => return Err(format!("flow:ap_zip => expected function from first stream, got {:?}", other)),
            };

            // Pull one value
            let x_val = apply_n_ary_function_value(interp, xs_clone, vec![]);
            let x_val = match x_val {
                Ok(v) => v,
                Err(e) => return Err(e),
            };

            // Apply
            apply_n_ary_function_value(interp, f_fn, vec![x_val])
        })),
        0,
    );

    Ok(Value::InkTransform(Box::new(clos)))
}

struct ApZipState {
    fs: Box<FunctionValue>,
    xs: Box<FunctionValue>,
}

// ---------------- ap (cartesian; buffers rhs) ----------------

fn flow_ap_bridge(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_cart_partial(None, None)),
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_cart_partial(None, None))
            } else {
                Ok(make_cart_partial(Some(a.clone()), None))
            }
        }
        2 => {
            let fs = &args[0];
            let xs = &args[1];
            let f_pl = is_placeholder(fs);
            let x_pl = is_placeholder(xs);
            if !f_pl && !x_pl {
                ap_finalize(fs.clone(), xs.clone())
            } else {
                Ok(make_cart_partial(
                    if f_pl { None } else { Some(fs.clone()) },
                    if x_pl { None } else { Some(xs.clone()) },
                ))
            }
        }
        n => Err(format!("flow:ap => expected up to 2 arguments, got {}", n)),
    }
}

fn make_cart_partial(fsrc_opt: Option<Value>, xsrc_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    Value::Function(Box::new(RustClosure(
        "flow:ap-partial".to_string(),
        Arc::new(Mutex::new(move |_interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut fsrc = fsrc_opt.clone();
            let mut xsrc = xsrc_opt.clone();

            for a in new_args {
                if fsrc.is_none() {
                    if !is_placeholder(&a) { fsrc = Some(a); }
                    continue;
                }
                if xsrc.is_none() {
                    if !is_placeholder(&a) { xsrc = Some(a); }
                    continue;
                }
                return Err("flow:ap => partial => too many arguments".to_string());
            }

            if fsrc.is_some() && xsrc.is_some() {
                ap_finalize(fsrc.unwrap(), xsrc.unwrap())
            } else {
                Ok(make_cart_partial(fsrc, xsrc))
            }
        })),
        0,
    )))
}

fn ap_finalize(fsrc: Value, xsrc: Value) -> Result<Value, String> {
    let fs_tf = to_transform(fsrc)?;
    let xs_tf = to_transform(xsrc)?;

    use mumu::parser::types::FunctionValue::RustClosure;

    let state = Arc::new(Mutex::new(ApState {
        fs: fs_tf,
        xs: xs_tf,
        xs_buf: None,
        cur_fn: None,
        xs_idx: 0,
        done: false,
    }));

    let clos = RustClosure(
        "flow:ap".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            let mut st = state.lock().map_err(|_| "flow:ap => state lock".to_string())?;

            if st.done {
                return Err("NO_MORE_DATA".to_string());
            }

            // Fill RHS buffer once
            if st.xs_buf.is_none() {
                let mut buf = Vec::new();
                loop {
                    let pull = apply_n_ary_function_value(interp, st.xs.clone(), vec![]);
                    match pull {
                        Ok(v) => buf.push(v),
                        Err(e) if e == "NO_MORE_DATA" => break,
                        Err(e) => return Err(e),
                    }
                }
                st.xs_buf = Some(buf);
                st.xs_idx = 0;
            }

            // If buffer empty, we're done.
            if st.xs_buf.as_ref().unwrap().is_empty() {
                st.done = true;
                return Err("NO_MORE_DATA".to_string());
            }

            loop {
                // If we don't have a current function, pull one
                if st.cur_fn.is_none() {
                    match apply_n_ary_function_value(interp, st.fs.clone(), vec![]) {
                        Ok(Value::Function(fb)) => {
                            st.cur_fn = Some(fb);
                            st.xs_idx = 0;
                        }
                        Ok(other) => return Err(format!("flow:ap => lhs must yield functions, got {:?}", other)),
                        Err(e) => {
                            st.done = true;
                            return Err(e);
                        }
                    }
                }

                // Emit next application without holding a long immutable borrow to the buffer.
                if let Some(fb) = st.cur_fn.clone() {
                    let xs_len = st.xs_buf.as_ref().unwrap().len();
                    if st.xs_idx < xs_len {
                        // Borrow only for the element clone, then release immediately.
                        let x = st.xs_buf.as_ref().unwrap()[st.xs_idx].clone();
                        st.xs_idx += 1;
                        let res = apply_n_ary_function_value(interp, fb, vec![x])?;
                        return Ok(res);
                    } else {
                        // exhausted xs for this function; pull next function
                        st.cur_fn = None;
                        continue;
                    }
                }
            }
        })),
        0,
    );

    Ok(Value::InkTransform(Box::new(clos)))
}

struct ApState {
    fs: Box<FunctionValue>,
    xs: Box<FunctionValue>,
    xs_buf: Option<Vec<Value>>,
    cur_fn: Option<Box<FunctionValue>>,
    xs_idx: usize,
    done: bool,
}