//! 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 [`lazy`].
//!
//! [`lazy`]: super::lazy

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

use crate::func::class_prelude::*;
use crate::func::derivations::{ApplicativeLA2ViaSeq, ApplicativeTupleViaLA2};

struct Wrapper<Q, F>(F, PhantomData<Q>);

trait Atom {
    fn next<'t>(self: Box<Self>) -> Oet<'t>
    where
        Self: 't;
}

fn atom<'a, F: 'a + FnOnce() -> Oet<'a>>(f: F) -> EvalTree<'a> {
    EvalTree::Atom(Box::new(Wrapper(f, PhantomData)))
}

fn batom<'a, F: 'a + FnOnce() -> Oet<'a>>(f: F) -> Box<EvalTree<'a>> {
    Box::new(atom(f))
}

fn satom<'a, F: 'a + FnOnce() -> Oet<'a>>(f: F) -> Oet<'a> {
    Some(atom(f))
}

impl<'a, F: 'a + FnOnce() -> Oet<'a>> Atom for Wrapper<&'a (), F> {
    fn next<'t>(self: Box<Self>) -> Oet<'t>
    where
        Self: 't,
    {
        (self.0)()
    }
}

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

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

impl<'a> EvalTree<'a> {
    fn next(self) -> Oet<'a> {
        match self {
            Self::Atom(f) => f.next(),
            Self::Composite(etl, etr) => match *etl {
                Self::Atom(f) => match f.next() {
                    Some(newleft) => Some(Self::Composite(Box::new(newleft), etr)),
                    None => Some(*etr),
                },
                Self::Composite(etll, etlr) => {
                    Some(Self::Composite(etll, Box::new(Self::Composite(etlr, etr))))
                }
            },
        }
    }
}

trait IntoOet<A> {
    fn into_oet<'t>(self: Box<Self>, f: Box<dyn 't + FnOnce(A)>) -> Oet<'t>
    where
        Self: 't,
        A: 't;
}

impl<'a, A: 'a, B: 'a, F: 'a + FnOnce(A) -> Stackless<'a, B>> IntoOet<B>
    for Wrapper<(&'a (), B), (Stackless<'a, A>, F)>
{
    fn into_oet<'t>(self: Box<Self>, f: Box<dyn 't + FnOnce(B)>) -> Oet<'t>
    where
        Self: 't,
        A: 't,
    {
        let cell_l = Rc::new(Cell::new(None));
        let cell_r = cell_l.clone();
        let (sstackless, sf) = self.0;
        Some(EvalTree::Composite(
            batom(move || sstackless.call(move |a| set_cell(cell_l, a))),
            batom(move || {
                let stackless = sf(get_cell(cell_r));
                satom(|| stackless.0.into_oet(f))
            }),
        ))
    }
}

impl<'a, A: 'a, B: 'a, F: 'a + FnOnce(A) -> B> IntoOet<B>
    for Wrapper<(&'a (), B, ()), (Stackless<'a, A>, F)>
{
    fn into_oet<'t>(self: Box<Self>, f: Box<dyn 't + FnOnce(B)>) -> Oet<'t>
    where
        Self: 't,
        A: 't,
    {
        let cell_l = Rc::new(Cell::new(None));
        let cell_r = cell_l.clone();
        let (sstackless, sf) = self.0;
        Some(EvalTree::Composite(
            batom(move || sstackless.call(move |a| set_cell(cell_l, a))),
            batom(move || {
                let b = sf(get_cell(cell_r));
                satom(|| {
                    f(b);
                    None
                })
            }),
        ))
    }
}

impl<A> IntoOet<A> for Wrapper<(), A> {
    fn into_oet<'t>(self: Box<Self>, f: Box<dyn 't + FnOnce(A)>) -> Oet<'t>
    where
        Self: 't,
        A: 't,
    {
        satom(|| {
            f(self.0);
            None
        })
    }
}

type StackessDyn<'a, A> = dyn 'a + IntoOet<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.into_oet(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(Wrapper((self, f), PhantomData)))
    }

    /// Method-like equivalent of [`Functor::fmap`].
    pub fn map<B: 'a>(self, f: impl 'a + FnOnce(A) -> B) -> Stackless<'a, B> {
        Stackless(Box::new(Wrapper((self, f), PhantomData)))
    }

    /// 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(Wrapper(value, PhantomData)))
    }
}

pub struct StacklessInstance;

impl WeakFunctorAny for StacklessInstance {
    type FAny<'a, A: 'a> = Stackless<'a, A>;
}

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

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

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

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

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

impl<'a> ApplicativeSelect<'a> for StacklessInstance {}

impl<'a> Applicative<'a> for StacklessInstance {}

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

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

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

    use super::StacklessInstance as T;

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

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