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

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

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

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

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

    Value::Function(Box::new(RustClosure(
        "flow:chain-partial".to_string(),
        Arc::new(Mutex::new(move |_interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut f = fn_opt.clone();
            let mut s = src_opt.clone();

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

            if f.is_some() && s.is_some() {
                chain_finalize(f.unwrap(), s.unwrap())
            } else {
                Ok(make_partial(f, s))
            }
        })),
        0,
    )))
}

fn chain_finalize(mapper: Box<FunctionValue>, src: Value) -> Result<Value, String> {
    let outer = to_transform(src)?;

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

    let state = Arc::new(Mutex::new(ChainState {
        mapper,
        outer,
        inner: None,
    }));

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

            loop {
                // If we have an inner stream, try to pull from it first
                if let Some(inner_tf) = st.inner.clone() {
                    match apply_n_ary_function_value(interp, inner_tf, vec![]) {
                        Ok(v) => return Ok(v),
                        Err(e) if e == "NO_MORE_DATA" => {
                            st.inner = None;
                            // continue to pull next from outer
                        }
                        Err(e) => return Err(e),
                    }
                }

                // Pull next from outer (a zero-arg transform function)
                let next_outer = apply_n_ary_function_value(interp, st.outer.clone(), vec![]);
                let item = match next_outer {
                    Ok(v) => v,
                    Err(e) if e == "NO_MORE_DATA" => return Err(e),
                    Err(e) => return Err(e),
                };

                // Map item -> stream/iterator/single value
                let mapped = apply_n_ary_function_value(interp, st.mapper.clone(), vec![item])?;
                st.inner = match mapped {
                    // If user returned a transform function directly, chain it
                    Value::Function(tf) => Some(tf),
                    // If user returned an Iterator handle, convert it to a transform
                    Value::Iterator(h) => Some(iter_to_transform(h)),
                    // Any other value: emit it once
                    other => Some(once_transform(other)),
                };
                // loop back to emit from inner immediately
            }
        })),
        0,
    );

    // Return the zero-arg transform function directly (no InkTransform wrapper)
    Ok(Value::Function(Box::new(clos)))
}

struct ChainState {
    mapper: Box<FunctionValue>,
    outer: Box<FunctionValue>,
    inner: Option<Box<FunctionValue>>,
}