//! Helper [Monad]s to move deep execution chains off the stack onto the heap.
//! [`Stackless<A>`] represents a wrapped value.
//!
//! For lazy stackful execution see [`super::lazy`].

use std::{cell::Cell, rc::Rc};

use crate::func::derivations::*;
use crate::func::*;

enum EvalTree<'a> {
    Atom(Box<dyn 'a + FnOnce() -> Oet<'a>>),
    Composite(Box<EvalTree<'a>>, Box<EvalTree<'a>>),
}

type Oet<'a> = Option<EvalTree<'a>>;

impl<'a> EvalTree<'a> {
    fn next(self) -> Oet<'a> {
        match self {
            EvalTree::Atom(f) => f(),
            EvalTree::Composite(left, right) => match *left {
                EvalTree::Atom(f) => match f() {
                    Some(newleft) => Some(EvalTree::Composite(Box::new(newleft), right)),
                    None => Some(*right),
                },
                EvalTree::Composite(leftleft, leftright) => Some(EvalTree::Composite(
                    leftleft,
                    Box::new(EvalTree::Composite(leftright, right)),
                )),
            },
        }
    }
}

type StackessDyn<'a, A> = dyn 'a + FnOnce(Box<dyn 'a + FnOnce(A)>) -> Oet<'a>;

pub struct Stackless<'a, A: 'a>(Box<StackessDyn<'a, A>>);

fn set_cell<A>(cell: Rc<Cell<Option<A>>>, a: A) {
    if cell.replace(Some(a)).is_some() {
        panic!("MITM overwritten")
    }
}

fn get_cell<A>(cell: Rc<Cell<Option<A>>>) -> A {
    match cell.replace(None) {
        Some(val) => val,
        None => panic!("MITM not set"),
    }
}

impl<'a, A: 'a> Stackless<'a, A> {
    fn call(self, f: impl 'a + FnOnce(A)) -> Oet<'a> {
        self.0(Box::new(f))
    }

    /// Method-like equivalent of [`Monad::bind`],
    /// the preferred way to chain [`Stackless<A>`] and `FnOnce(A) -> Stackless<B>` into [`Stackless<B>`].
    pub fn bind<B: 'a>(self, f: impl 'a + FnOnce(A) -> Stackless<'a, B>) -> Stackless<'a, B> {
        Stackless(Box::new(|takesb| {
            let lcell = Rc::new(Cell::new(None));
            let rcell = lcell.clone();
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(move || {
                    self.call(move |a| set_cell(lcell, a))
                }))),
                Box::new(EvalTree::Atom(Box::new(move || {
                    let stackless = f(get_cell(rcell));
                    Some(EvalTree::Atom(Box::new(|| stackless.0(takesb))))
                }))),
            ))
        }))
    }

    /// Method-like equivalent of [`Functor::fmap`].
    pub fn map<B: 'a>(self, f: impl 'a + FnOnce(A) -> B) -> Stackless<'a, B> {
        Stackless(Box::new(|takesb| {
            let lcell = Rc::new(Cell::new(None));
            let rcell = lcell.clone();
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(move || {
                    self.call(move |a| set_cell(lcell, a))
                }))),
                Box::new(EvalTree::Atom(Box::new(move || {
                    let b = f(get_cell(rcell));
                    Some(EvalTree::Atom(Box::new(|| {
                        takesb(b);
                        None
                    })))
                }))),
            ))
        }))
    }

    /// Evaluate. Process is loop-like on the inside
    /// with the least amount of recursion the current model allows to use.
    pub fn evaluate(self) -> A {
        let ocell = Rc::new(Cell::new(None));
        let icell = ocell.clone();
        let mut eval = self.call(|a| set_cell(icell, a));
        while let Some(tree) = eval {
            eval = tree.next()
        }
        get_cell(ocell)
    }
}

impl<'a, A: 'a> From<A> for Stackless<'a, A> {
    fn from(value: A) -> Self {
        Stackless(Box::new(|takesa| {
            Some(EvalTree::Atom(Box::new(|| {
                takesa(value);
                None
            })))
        }))
    }
}

pub struct StacklessClass;

impl WeakFunctor for StacklessClass {
    type F<'a, A: 'a> = Stackless<'a, A>;
}

impl Functor for StacklessClass {
    fn fmap<'a, A: 'a, B: 'a>(f: impl 'a + FnOnce(A) -> B, fa: Self::F<'a, A>) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        fa.map(f)
    }

    fn replace<'a, A: 'a, B: 'a>(fa: Self::F<'a, A>, b: B) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        Stackless(Box::new(|takesb| {
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(move || fa.call(drop)))),
                Box::new(EvalTree::Atom(Box::new(move || {
                    takesb(b);
                    None
                }))),
            ))
        }))
    }
}

impl Pure for StacklessClass {
    fn pure<'a, A: 'a>(a: A) -> Self::F<'a, A> {
        Stackless::from(a)
    }
}

impl ApplicativeSeq for StacklessClass {
    fn seq<'a, A: 'a, B: 'a>(
        ff: Self::F<'a, impl 'a + FnOnce(A) -> B>,
        fa: Self::F<'a, A>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        ff.bind(|f| fa.map(f))
    }
}

impl ApplicativeLA2 for StacklessClass {
    fn la2<'a, A: 'a, B: 'a, C: 'a>(
        f: impl 'a + FnOnce(A, B) -> C,
        fa: Self::F<'a, A>,
        fb: Self::F<'a, B>,
    ) -> Self::F<'a, C>
    where
        Self: 'a,
    {
        Self::_la2_via_seq(f, fa, fb)
    }
}

impl ApplicativeTuple for StacklessClass {
    fn tuple<'a, A: 'a, B: 'a>((fa, fb): (Self::F<'a, A>, Self::F<'a, B>)) -> Self::F<'a, (A, B)>
    where
        Self: 'a,
    {
        Self::_tuple_via_la2((fa, fb))
    }
}

impl Applicative for StacklessClass {
    fn discard_first<'a, A: 'a, B: 'a>(fa: Self::F<'a, A>, fb: Self::F<'a, B>) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        Stackless(Box::new(|takesb| {
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(|| fa.call(drop)))),
                Box::new(EvalTree::Atom(Box::new(|| fb.0(takesb)))),
            ))
        }))
    }

    fn discard_second<'a, A: 'a, B: 'a>(fa: Self::F<'a, A>, fb: Self::F<'a, B>) -> Self::F<'a, A>
    where
        Self: 'a,
    {
        Stackless(Box::new(|takesa| {
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(|| fa.0(takesa)))),
                Box::new(EvalTree::Atom(Box::new(|| fb.call(drop)))),
            ))
        }))
    }
}

impl Monad for StacklessClass {
    fn bind<'a, A: 'a, B: 'a>(
        fa: Self::F<'a, A>,
        f: impl 'a + FnOnce(A) -> Self::F<'a, B>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        fa.bind(f)
    }

    fn iterate_mut<'a, A: 'a, B: 'a>(
        a: A,
        mut f: impl 'a + FnMut(A) -> Self::F<'a, ControlFlow<B, A>>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        Self::pure(a).bind(move |a| {
            f(a).bind(|state| match state {
                ControlFlow::Continue(next_a) => Self::iterate_mut(next_a, f),
                ControlFlow::Break(b) => Self::pure(b),
            })
        })
    }

    fn iterate_argument<'a, A: 'a, B: 'a>(
        a: A,
        f: impl AIterative<'a, T = Self, A = A, B = B>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        Self::pure(a).bind(move |a| {
            f.next(a).bind(|state| match state {
                ControlFlow::Continue((next_a, next_f)) => Self::iterate_argument(next_a, next_f),
                ControlFlow::Break(b) => Self::pure(b),
            })
        })
    }

    fn iterate<'a, B: 'a>(f: impl Iterative<'a, T = Self, B = B>) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        Self::pure(()).bind(move |_| {
            f.next().bind(|state| match state {
                ControlFlow::Continue(next_f) => Self::iterate(next_f),
                ControlFlow::Break(b) => Self::pure(b),
            })
        })
    }

    fn join<'a, A: 'a>(ffa: Self::F<'a, Self::F<'a, A>>) -> Self::F<'a, A>
    where
        Self::F<'a, A>: 'a,
        Self: 'a,
    {
        Stackless(Box::new(|takesa| {
            let lcell = Rc::new(Cell::new(None));
            let rcell = lcell.clone();
            Some(EvalTree::Composite(
                Box::new(EvalTree::Atom(Box::new(move || {
                    ffa.call(move |a| set_cell(lcell, a))
                }))),
                Box::new(EvalTree::Atom(Box::new(move || {
                    let stackless = get_cell(rcell);
                    Some(EvalTree::Atom(Box::new(|| stackless.0(takesa))))
                }))),
            ))
        }))
    }
}

#[cfg(test)]
mod stackless_test {
    use super::{test_suite, tests, Stackless};

    use super::StacklessClass as T;

    impl tests::Eqr for T {
        fn eqr<'a, A: PartialEq + std::fmt::Debug + 'a>(
            name: &'a str,
            left: Self::F<'a, A>,
            right: Self::F<'a, A>,
        ) -> tests::R {
            tests::eqr(name, left.evaluate(), right.evaluate())
        }
    }

    impl test_suite::FunctorTestSuite for T {
        fn sample<'a, A: 'a, F: FnMut(&'a dyn Fn(A) -> Self::F<'a, A>)>(mut f: F)
        where
            Self::F<'a, A>: 'a,
        {
            f(&|a| a.into());
        }
    }

    #[test]
    fn monad_follows_laws() {
        test_suite::monad_follows_laws::<T>().unwrap();
    }

    fn factorial(n: u32) -> Stackless<'static, u32> {
        if n > 0 {
            Stackless::from(()).bind(move |_| factorial(n - 1).map(move |acc| acc * n))
        } else {
            1.into()
        }
    }

    fn dumb(n: u32) -> Stackless<'static, u32> {
        if n > 0 {
            Stackless::from(()).bind(move |_| dumb(n - 1).map(move |acc| acc + 1))
        } else {
            0.into()
        }
    }

    #[ignore]
    #[test]
    fn test_factorial() {
        assert_eq!(factorial(10).evaluate(), 3628800);
    }

    #[ignore]
    #[test]
    fn test_dumb() {
        let n = 1000;
        assert_eq!(dumb(n).evaluate(), n);
    }
}