flow-mumu 0.2.0-rc.9

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

use mumu::parser::interpreter::{DynamicFnInfo, Interpreter};
use mumu::parser::interpreter::apply::apply_n_ary_function_value;
use mumu::parser::types::{
    FunctionValue, IteratorHandle, IteratorKind, Value,
};

/// Register `flow:trans` (now with **eager** evaluation when fully-applied)
pub fn register_flow_trans(interp: &mut Interpreter) {
    let bridging = Arc::new(Mutex::new(flow_trans_bridge_fn));
    let info = DynamicFnInfo::new(bridging, true);
    interp.register_dynamic_function_ex("flow:trans", info);

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

/// Alias: map ≡ trans (same eager behavior on 2-arg calls)
pub fn register_flow_map_alias(interp: &mut Interpreter) {
    let bridging = Arc::new(Mutex::new(flow_trans_bridge_fn));
    let info = DynamicFnInfo::new(bridging, true);
    interp.register_dynamic_function_ex("flow:map", info);

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

/// Bridge:
/// - 0 args → partial (waiting for fn and source)
/// - 1 arg  → partial
/// - 2 args → **EAGER** map now, returning a concrete array (no RustClosure crosses the boundary)
fn flow_trans_bridge_fn(intp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    if intp.is_verbose() {
        eprintln!("[flow:trans] bridge args={:?}", short_val_list(&args));
    }

    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:trans => first arg must be Function or '_', got {:?}",
                    a
                ))
            }
        }
        2 => {
            let first = &args[0];
            let second = &args[1];

            // Fully-specified – go eager.
            let func = match first {
                Value::Function(fb) => fb.clone(),
                other => {
                    return Err(format!(
                        "flow:trans => first param must be a Function, got {:?}",
                        other
                    ))
                }
            };

            match second {
                Value::Iterator(h) => {
                    if intp.is_verbose() {
                        eprintln!("[flow:trans] eager over Iterator");
                    }
                    eager_map_over_iterator(intp, func, h.clone())
                }
                Value::Function(up_tf) => {
                    if intp.is_verbose() {
                        eprintln!("[flow:trans] eager over upstream transform (zero-arg function)");
                    }
                    eager_map_over_transform(intp, func, up_tf.clone())
                }
                other => Err(format!(
                    "flow:trans => second param must be Iterator or transform Function, got {:?}",
                    other
                )),
            }
        }
        n => Err(format!("flow:trans => expected up to 2 arguments, got {}", n)),
    }
}

/// Keep the partial in case user supplies `_` or does curried usage.
/// This remains a thin shell; only the fully-applied (2-arg) path is eager.
fn make_partial(func_opt: Option<Box<FunctionValue>>, source_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let st = TransPartialState {
        func: func_opt,
        source: source_opt,
    };

    Value::Function(Box::new(RustClosure(
        "flow:trans-partial".to_string(),
        Arc::new(Mutex::new(move |intp: &mut Interpreter, new_args: Vec<Value>| {
            let mut f_cur = st.func.clone();
            let mut s_cur = st.source.clone();

            for arg in &new_args {
                if f_cur.is_none() {
                    if is_placeholder(arg) {
                        // remain None
                    } else if let Value::Function(fb) = arg {
                        f_cur = Some(fb.clone());
                    } else {
                        return Err(format!(
                            "flow:trans => first param must be Function or '_', got {:?}",
                            arg
                        ));
                    }
                    continue;
                }
                if s_cur.is_none() {
                    if is_placeholder(arg) {
                        // remain None
                    } else {
                        s_cur = Some(arg.clone());
                    }
                    continue;
                }
                return Err("flow:trans => partial => too many arguments".to_string());
            }

            // If both present now, run the same eager path as the 2-arg bridge
            if let (Some(fb), Some(src)) = (f_cur.clone(), s_cur.clone()) {
                match src {
                    Value::Iterator(h) => eager_map_over_iterator(intp, fb, h),
                    Value::Function(tf) => eager_map_over_transform(intp, fb, tf),
                    other => Err(format!(
                        "flow:trans => second param must be Iterator or transform Function, got {:?}",
                        other
                    )),
                }
            } else {
                Ok(make_partial(f_cur, s_cur))
            }
        })),
        0,
    )))
}

/* ────────────────────────────────────────────────────────────────────────────
   EAGER implementations
   ───────────────────────────────────────────────────────────────────────── */

fn eager_map_over_iterator(
    intp: &mut Interpreter,
    mapper: Box<FunctionValue>,
    iter: IteratorHandle,
) -> Result<Value, String> {
    let mut out: Vec<Value> = Vec::new();

    match &iter.kind {
        IteratorKind::Core(state_arc) => {
            loop {
                let item_opt = {
                    let mut g = state_arc
                        .lock()
                        .map_err(|_| "flow:trans => iterator lock error".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)
                    }
                };
                let Some(item) = item_opt else { break };

                let mapped = apply_n_ary_function_value(intp, mapper.clone(), vec![item])?;
                out.push(mapped);
            }
        }
        IteratorKind::Plugin(plugin_arc) => {
            let mut plugin = plugin_arc
                .lock()
                .map_err(|_| "flow:trans => plugin iterator lock error".to_string())?;
            loop {
                match plugin.next_value() {
                    Ok(item) => {
                        let mapped = apply_n_ary_function_value(intp, mapper.clone(), vec![item])?;
                        out.push(mapped);
                    }
                    Err(e) if e == "NO_MORE_DATA" => break,
                    Err(e) => return Err(e),
                }
            }
        }
    }

    let folded = fold_to_homogeneous_array(out);
    if intp.is_verbose() {
        eprintln!("[flow:trans] eager Iterator => {}", short_val(&folded));
    }
    Ok(folded)
}

fn eager_map_over_transform(
    intp: &mut Interpreter,
    mapper: Box<FunctionValue>,
    upstream_tf: Box<FunctionValue>,
) -> Result<Value, String> {
    let mut out: Vec<Value> = Vec::new();

    loop {
        match apply_n_ary_function_value(intp, upstream_tf.clone(), vec![]) {
            Ok(item) => {
                let mapped = apply_n_ary_function_value(intp, mapper.clone(), vec![item])?;
                out.push(mapped);
            }
            Err(e) if e == "NO_MORE_DATA" => break,
            Err(e) => return Err(e),
        }
    }

    let folded = fold_to_homogeneous_array(out);
    if intp.is_verbose() {
        eprintln!("[flow:trans] eager Transform => {}", short_val(&folded));
    }
    Ok(folded)
}

/// Convert a vector of Values into the tightest homogeneous array,
/// or MixedArray if heterogeneous.
fn fold_to_homogeneous_array(mut items: Vec<Value>) -> Value {
    if items.is_empty() {
        return Value::IntArray(Vec::new());
    }

    let all_int = items.iter().all(|v| matches!(v, Value::Int(_)));
    if all_int {
        let mut out = Vec::with_capacity(items.len());
        for v in items.drain(..) {
            if let Value::Int(i) = v {
                out.push(i);
            }
        }
        return Value::IntArray(out);
    }

    let all_float = items.iter().all(|v| matches!(v, Value::Float(_)));
    if all_float {
        let mut out = Vec::with_capacity(items.len());
        for v in items.drain(..) {
            if let Value::Float(f) = v {
                out.push(f);
            }
        }
        return Value::FloatArray(out);
    }

    let all_bool = items.iter().all(|v| matches!(v, Value::Bool(_)));
    if all_bool {
        let mut out = Vec::with_capacity(items.len());
        for v in items.drain(..) {
            if let Value::Bool(b) = v {
                out.push(b);
            }
        }
        return Value::BoolArray(out);
    }

    let all_str = items.iter().all(|v| matches!(v, Value::SingleString(_)));
    if all_str {
        let mut out = Vec::with_capacity(items.len());
        for v in items.drain(..) {
            if let Value::SingleString(s) = v {
                out.push(s);
            }
        }
        return Value::StrArray(out);
    }

    Value::MixedArray(items)
}

/* ────────────────────────────────────────────────────────────────────────────
   Small helpers
   ───────────────────────────────────────────────────────────────────────── */

#[derive(Clone)]
struct TransPartialState {
    func: Option<Box<FunctionValue>>,
    source: Option<Value>,
}

fn is_placeholder(val: &Value) -> bool {
    match val {
        Value::Placeholder => true,
        Value::SingleString(s) if s == "_" => true,
        Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_" => true,
        _ => false,
    }
}

fn short_val(v: &Value) -> String {
    match v {
        Value::IntArray(v) => format!("[int; {}]", v.len()),
        Value::FloatArray(v) => format!("[float; {}]", v.len()),
        Value::BoolArray(v) => format!("[bool; {}]", v.len()),
        Value::StrArray(v) => format!("[str; {}]", v.len()),
        Value::MixedArray(v) => format!("[mixed; {}]", v.len()),
        Value::Function(_) => "[Function]".to_string(),
        Value::Iterator(_) => "[Iterator]".to_string(),
        other => format!("{:?}", other),
    }
}

fn short_val_list(xs: &[Value]) -> String {
    let parts: Vec<String> = xs.iter().map(short_val).collect();
    format!("[{}]", parts.join(", "))
}