mumu-flow 0.1.0

use std::sync::{Arc, Mutex};
use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::interpreter::apply_n_ary_function_value;
use mumu::parser::types::{Value, FunctionValue, InkIteratorKind};

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

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

fn flow_slice_bridge_fn(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    make_slice_partial(args)
}

fn make_slice_partial(existing_args: Vec<Value>) -> Result<Value, String> {
    use mumu::parser::types::FunctionValue::RustClosure;

    let store = Arc::new(SlicePartialState { arg_list: existing_args });

    let clos = RustClosure(
        "flow:slice-partial".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut combined = store.arg_list.clone();
            for na in &new_args {
                combined.push(na.clone());
            }
            finalize_or_partial(interp, combined)
        })),
        0
    );

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

fn finalize_or_partial(_interp: &mut Interpreter, combined: Vec<Value>) -> Result<Value, String> {
    if combined.len() < 3 {
        return make_slice_partial(combined);
    }
    if combined.len() > 3 {
        return Err(format!("flow:slice => expected up to 3 arguments, got {}", combined.len()));
    }

    let skip_val = &combined[0];
    let count_val = &combined[1];
    let source_val = &combined[2];

    if is_placeholder(skip_val) || is_placeholder(count_val) || is_placeholder(source_val) {
        return make_slice_partial(combined);
    }

    let skip_int = match skip_val {
        Value::Int(i) => *i,
        _ => return Err(format!("flow:slice => 'skip' must be int or placeholder, got {:?}", skip_val)),
    };
    let count_int = match count_val {
        Value::Int(i) => *i,
        _ => return Err(format!("flow:slice => 'count' must be int or placeholder, got {:?}", count_val)),
    };

    match source_val {
        Value::InkIterator(h) => Ok(Value::InkTransform(build_slice_transform(skip_int, count_int, SourceKind::InkIterator(h.clone())))),
        Value::InkTransform(fb) => Ok(Value::InkTransform(build_slice_transform(skip_int, count_int, SourceKind::InkTransform(fb.clone())))),
        _ => Err(format!("flow:slice => third arg must be InkIterator or InkTransform, got {:?}", source_val)),
    }
}

enum SourceKind {
    InkIterator(mumu::parser::types::InkIteratorHandle),
    InkTransform(Box<FunctionValue>),
}

fn build_slice_transform(skip: i32, count: i32, source: SourceKind) -> Box<FunctionValue> {
    use mumu::parser::types::FunctionValue::RustClosure;

    let data = Arc::new(Mutex::new(SliceData {
        skip_so_far: 0,
        taken_so_far: 0,
        skip_limit: skip.max(0),
        take_limit: count.max(0),
        source_kind: source,
        done: false,
    }));

    Box::new(RustClosure(
        "flow:slice-transform".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            let mut locked = data.lock().map_err(|_| "flow:slice => lock error".to_string())?;
            if locked.done || locked.taken_so_far >= locked.take_limit {
                return Err("NO_MORE_DATA".to_string());
            }

            while locked.skip_so_far < locked.skip_limit {
                let skip_result = {
                    match &mut locked.source_kind {
                        SourceKind::InkIterator(handle) => {
                            match &mut handle.kind {
                                InkIteratorKind::Core(state_arc) => {
                                    let mut g = state_arc.lock().map_err(|_| "flow:slice => InkIterator lock error".to_string())?;
                                    if g.done || g.current >= g.end {
                                        true
                                    } else {
                                        g.current += 1;
                                        false
                                    }
                                }
                                InkIteratorKind::Plugin(_) => {
                                    true // or false? adjust if plugin needed
                                }
                            }
                        }
                        SourceKind::InkTransform(fb) => {
                            match apply_n_ary_function_value(interp, fb.clone(), vec![]) {
                                Ok(_) => false,
                                Err(e) if e == "NO_MORE_DATA" => true,
                                Err(e) => return Err(e),
                            }
                        }
                    }
                };
                if skip_result {
                    locked.done = true;
                    return Err("NO_MORE_DATA".to_string());
                }
                locked.skip_so_far += 1;
                if locked.skip_so_far >= locked.skip_limit {
                    break;
                }
            }

            let next_val = {
                match &mut locked.source_kind {
                    SourceKind::InkIterator(handle) => {
                        match &mut handle.kind {
                            InkIteratorKind::Core(state_arc) => {
                                let mut g = state_arc.lock().map_err(|_| "flow:slice => InkIterator lock error".to_string())?;
                                if g.done || g.current >= g.end {
                                    None
                                } else {
                                    let cur_item = g.current;
                                    g.current += 1;
                                    Some(Value::Int(cur_item as i32))
                                }
                            }
                            InkIteratorKind::Plugin(_) => None
                        }
                    }
                    SourceKind::InkTransform(fb) => {
                        match apply_n_ary_function_value(interp, fb.clone(), vec![]) {
                            Ok(v) => Some(v),
                            Err(e) if e == "NO_MORE_DATA" => None,
                            Err(e) => return Err(e),
                        }
                    }
                }
            };

            match next_val {
                None => {
                    locked.done = true;
                    Err("NO_MORE_DATA".to_string())
                }
                Some(val) => {
                    locked.taken_so_far += 1;
                    if locked.taken_so_far >= locked.take_limit {
                        locked.done = true;
                    }
                    Ok(val)
                }
            }
        })),
        0
    ))
}

struct SliceData {
    skip_so_far: i32,
    taken_so_far: i32,
    skip_limit: i32,
    take_limit: i32,
    source_kind: SourceKind,
    done: bool,
}

struct SlicePartialState {
    arg_list: Vec<Value>,
}

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