use std::{cell::RefCell, rc::Rc};
use matrix_lang::prelude::*;
use matrix_macros::native_func;
use crate::{error, next, VmArgs, unpack_args, unpack_varargs};

use Value as V;

#[native_func(1)]
fn len(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [value] = unpack_args!(args);
    let len = match value {
        V::List(l) => l.len(),
        V::Matrix(m) => m.values.len(),
        V::String(s) => s.len(),
        V::Table(t) => t.len(),
        _ => 1
    };
    Ok(V::Int(len as i64))
}

#[native_func(1)]
fn sum((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut res = next!(vm, frame, iter)?;
    if res == Value::Nil {
        return Ok(res)
    };
    loop {
        let next = next!(vm, frame, iter)?;
        if next == Value::Nil { break }
        res = (res + next)?;
    }
    Ok(res)
}

#[native_func(0..)]
fn range(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let ([], varargs) = unpack_varargs!(args);
    let (min, max, step) = match varargs.as_slice() {
        [V::Int(max)] => {
            (0, *max, 1)
        },
        [V::Int(min), V::Int(max)] => {
            (*min, *max, 1)
        },
        [V::Int(min), V::Int(max), V::Int(step)] => {
            (*min, *max, *step)
        },
        _ => return error!("range takes 1 to 3 [Int]'s")
    };

    let i = RefCell::new(min);
    Ok(iter!(move |_,_| {
        let curr = *(i.borrow());
        *(i.borrow_mut()) = curr + step;
        if (step >= 0 && curr >= max) || (step <= 0 && curr <= max) {
            return Ok(V::Nil)
        } else {
            return Ok(V::Int(curr))
        }
    }))
}

#[native_func(1)]
fn iter(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [value] = unpack_args!(args);
    Ok(value.into_iter()?)
}

#[native_func(1)]
fn once(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [value] = unpack_args!(args);
    let first = RefCell::new(false);
    Ok(iter!(move |_,_| {
        if *first.borrow() == false {
            *first.borrow_mut() = true;
            Ok(value.clone())
        } else {
            Ok(Value::Nil)
        }
    }))
}

#[native_func(1)]
fn list((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut list = Vec::new();
    loop {
        let res = next!(vm, frame, iter)?;
        if res == Value::Nil { break }
        list.push(res);
    }
    Ok(Value::List(list.into()))
}

#[native_func(2)]
fn map(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [fun, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Function(fun) = fun else {
        return error!("map 1st arg must be a function")
    };
    Ok(iter!(move |(vm,frame),_| {
        let input = next!(vm, frame, iter)?;
        if input == Value::Nil {
            return Ok(input)
        }
        vm.run_fn(frame, fun.clone(), vec![input])
    }))
}

#[native_func(2)]
fn fold((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [fun, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Function(fun) = fun else {
        return error!("fold 1st arg must be a function")
    };
    let mut res = next!(vm, frame, iter)?;
    if res == Value::Nil {
        return Ok(Value::Nil)
    };
    loop {
        let next = next!(vm, frame, iter)?;
        if next == Value::Nil {
            return Ok(res)
        };
        res = vm.run_fn(frame, fun.clone(), vec![res, next])?;
    }
}

#[native_func(3)]
fn foldi((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [init, fun, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Function(fun) = fun else {
        return error!("foldi 1st arg must be a function")
    };
    let mut res = init;
    loop {
        let next = next!(vm, frame, iter)?;
        if next == Value::Nil {
            return Ok(res)
        };
        res = vm.run_fn(frame, fun.clone(), vec![res, next])?;
    }
}

#[native_func(1)]
fn count((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut len = 0;
    loop {
        let res = vm.run_fn(frame, iter.clone(), vec![])?;
        if res == Value::Nil { break };
        len += 1;
    }
    Ok(Value::Int(len))
}

#[native_func(3)]
fn scan(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [init, fun, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Function(fun) = fun else {
        return error!("scan 2nd arg must be a function")
    };
    let res = RefCell::new(init);
    Ok(iter!(move |(vm,frame),_| {
        let next = next!(vm, frame, iter)?;
        if next == Value::Nil {
            return Ok(Value::Nil)
        };
        let res_next = vm.run_fn(frame, fun.clone(), vec![res.borrow().clone(), next])?;
        *res.borrow_mut() = res_next;
        Ok(res.borrow().clone())
    }))
}

#[native_func(2)]
fn filter(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [fun, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Function(fun) = fun else {
        return error!("filter 1st arg must be a function")
    };
    Ok(iter!(move |(vm,frame),_| {
        loop {
            let next = next!(vm, frame, iter)?;
            if next == Value::Nil {
                return Ok(Value::Nil)
            }
            let res = vm.run_fn(frame, fun.clone(), vec![next.clone()])?;
            if !!res {
                return Ok(next)
            }
        }
    }))
}

#[native_func(0..)]
fn chain(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let ([], iters) = unpack_varargs!(args);

    let mut chaind = Vec::new();
    for iter in iters {
        chaind.push(iter.into_iter_fn()?);
    }

    chaind.reverse();
    let chaind = RefCell::new(chaind);
    Ok(iter!(move |(vm,frame), _| {
        loop {
            let curr = match chaind.borrow_mut().last() {
                Some(iter) => iter.clone(),
                None => return Ok(Value::Nil)
            };
            match vm.run_fn(frame, curr.clone(), vec![]) {
                Ok(Value::Nil) => {
                    chaind.borrow_mut().pop();
                    continue;
                },
                v => return v
            };
        }
    }))
}

#[native_func(1)]
fn lines(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [str] = unpack_args!(args);
    let Value::String(str) = str else {
        return error!("lines arg must be a string")
    };
    let lines: Vec<Rc<str>> = str.split("\n").map(|s| Rc::from(s)).collect();
    let res = RefCell::new(lines.into_iter());
    Ok(iter!(move |_,_| {
        match res.borrow_mut().next() {
            Some(line) => Ok(Value::String(line)),
            None => Ok(Value::Nil)
        }
    }))
}

#[native_func(2)]
fn skip((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [count, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Int(i) = count else {
        return error!("skip count requires a int")
    };
    for _ in 0..i {
        next!(vm, frame, iter)?;
    }
    Ok(Value::Iter(iter))
}

#[native_func(2)]
fn alternate(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [il, ir] = unpack_args!(args);
    let il = il.into_iter_fn()?;
    let ir = ir.into_iter_fn()?;
    let flag = RefCell::new(Some(0));
    Ok(iter!(move |(vm, frame),_| {
        let f = *flag.borrow();
        match f {
            Some(0) => {
                let val = next!(vm, frame, il)?;
                if val == Value::Nil {
                    *flag.borrow_mut() = None;
                } else {
                    *flag.borrow_mut() = Some(1);
                }
                Ok(val)
            },
            Some(1) => {
                let val = next!(vm, frame, ir)?;
                if val == Value::Nil {
                    *flag.borrow_mut() = None;
                } else {
                    *flag.borrow_mut() = Some(0);
                }
                Ok(val)
            },
            _ => Ok(Value::Nil)
        }
    }))
}

#[native_func(2)]
fn intersperse((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [value, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let flag = RefCell::new(Some(0));
    let next = RefCell::new(next!(vm, frame, iter)?);
    Ok(iter!(move |(vm, frame),_| {
        let f = *flag.borrow();
        match f {
            Some(0) => {
                let val = next.borrow();
                if *val == Value::Nil {
                    *flag.borrow_mut() = None;
                } else {
                    *flag.borrow_mut() = Some(1);
                }
                Ok(val.clone())
            },
            Some(1) => {
                *next.borrow_mut() = next!(vm, frame, iter)?;
                if *next.borrow() == Value::Nil {
                    *flag.borrow_mut() = None;
                    return Ok(Value::Nil)
                } else {
                    *flag.borrow_mut() = Some(0);
                }
                let val = value.clone();
                Ok(val)
            },
            _ => Ok(Value::Nil)
        }
    }))
}

#[native_func(2)]
fn zip(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [il, ir] = unpack_args!(args);
    let il = il.into_iter_fn()?;
    let ir = ir.into_iter_fn()?;
    let flag = RefCell::new(true);
    Ok(iter!(move |(vm, frame),_| {
        let f = *flag.borrow();
        match f {
            true => {
                let vl = next!(vm, frame, il)?;
                let vr = next!(vm, frame, ir)?;
                if vl == Value::Nil || vr == Value::Nil {
                    *flag.borrow_mut() = false;
                    return Ok(Value::Nil)
                }
                Ok(Value::List(vec![vl, vr].into()))
            },
            _ => Ok(Value::Nil)
        }
    }))
}

#[native_func(1)]
fn unzip((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut ll = Vec::new();
    let mut lr = Vec::new();
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil {
            break;
        }
        let Value::List(vals) = val else {
            return error!("unzip only works over a iterator of pairs");
        };
        let vals = vals.into_inner();
        if vals.len() != 2 {
            return error!("unzip only works over a iterator of pairs");
        }
        let [l, r] = vals.try_into().map_err(|_| exception!(RUNTIME_EXCEPTION, "can only unzip over a set of pairs"))?;
        ll.push(l);
        lr.push(r);
    }
    let ll = Value::List(ll.into());
    let lr = Value::List(lr.into());
    Ok(Value::List(vec![ll, lr].into()))
}

#[native_func(1)]
fn cycle((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut values = Vec::new();
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        values.push(val);
    }
    let idx = RefCell::new(0_usize);
    Ok(iter!(move |_,_| {
        let i = *idx.borrow();
        *idx.borrow_mut() += 1;
        *idx.borrow_mut() %= values.len();
        Ok(values[i].clone())
    }))
}

#[native_func(2)]
fn take((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [count, iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let Value::Int(count) = count else {
        return error!("take requires an int amount to collect")
    };
    let mut values = Vec::new();
    for _ in 0..count {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        values.push(val);
    }
    Ok(Value::List(values.into()))
}

#[native_func(1)]
fn last((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut last = Value::Nil;
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        last = val;
    }
    Ok(last)
}

#[native_func(1)]
fn next((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let val = next!(vm, frame, iter)?;
    Ok(val)
}

#[native_func(1)]
fn min((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut min = Value::Nil;
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        if min == Value::Nil || val < min {
            min = val;
        }
    }
    Ok(min)
}

#[native_func(1)]
fn max((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut max = Value::Nil;
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        if max == Value::Nil || val > max {
            max = val;
        }
    }
    Ok(max)
}

#[native_func(1)]
fn rev((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;
    let mut values = Vec::new();
    loop {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil { break };
        values.push(val);
    }
    let values = RefCell::new(values);
    Ok(iter!(move |_,_| {
        match values.borrow_mut().pop() {
            Some(v) => Ok(v),
            None => Ok(Value::Nil)
        }
    }))
}

#[native_func(1)]
fn enumerate(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [iter] = unpack_args!(args);
    let iter = iter.into_iter_fn()?;

    let idx = RefCell::new(0_i64);
    Ok(iter!(move |(vm, frame),_| {
        let val = next!(vm, frame, iter)?;
        if val == Value::Nil {
            return Ok(Value::Nil)
        };
        let curr = *idx.borrow();
        *idx.borrow_mut() += 1;
        Ok(Value::List(vec![Value::Int(curr), val].into()))
    }))
}

#[native_func(1)]
fn iterable(_: VmArgs, args: Vec<Value>) -> Result<Value> {
    let [value] = unpack_args!(args);
    use Value as V;
    match value {
        V::Function(_) |
            V::List(_) |
            V::Range(_) |
            V::Iter(_)
            => Ok(V::Bool(true)),
        _ => Ok(V::Bool(false))
    }
}

pub fn load(vm: &mut Vm) {
    vm.load_global_fn(take(), "take");
    vm.load_global_fn(unzip(), "unzip");
    vm.load_global_fn(count(), "count");
    vm.load_global_fn(len(), "len");
    vm.load_global_fn(sum(), "sum");
    vm.load_global_fn(min(), "min");
    vm.load_global_fn(max(), "max");
    vm.load_global_fn(next(), "next");
    vm.load_global_fn(last(), "last");

    vm.load_global_fn(list(), "list");
    vm.load_global_fn(fold(), "fold");
    vm.load_global_fn(foldi(), "foldi");
    vm.load_global_fn(scan(), "scan");
    vm.load_global_fn(chain(), "chain");
    vm.load_global_fn(lines(), "lines");
    vm.load_global_fn(skip(), "skip");

    vm.load_global_fn(once(), "once");
    vm.load_global_fn(iter(), "iter");
    vm.load_global_fn(range(), "range");
    vm.load_global_fn(map(), "map");
    vm.load_global_fn(filter(), "filter");
    vm.load_global_fn(skip(), "skip");
    vm.load_global_fn(zip(), "zip");
    vm.load_global_fn(cycle(), "cycle");
    vm.load_global_fn(alternate(), "alternate");
    vm.load_global_fn(intersperse(), "intersperse");
    vm.load_global_fn(rev(), "rev");
    vm.load_global_fn(enumerate(), "enumerate");

    vm.load_global_fn(iterable(), "iterable");
}