//! Structures for tracing the execution flow of [Monad]s.

mod rendered;
mod rendered_display;
mod trace;
mod traceable;
mod traced;

use crate::core::*;

use super::*;

pub use self::rendered::*;
use self::trace::*;
pub use self::traceable::*;
pub use self::traced::*;

/// [`Diagnostic`] for [Traced] objects.
///
/// [`Diagnostic::after`]/[`Diagnostic::before`] are represented in [`RenderedCommon::Event`].
/// [`Diagnostic::wrapped`] is represented in [`RenderedCommon::Action`].
pub struct TracedDiagnostic;

/// Implementation of [`Monad`] for [Traced] objects.
pub struct TracedClass;

trait WithTrace: Sized {
    fn with_trace(self, t: TraceBox) -> Traced<Self>;
}

impl<A> WithTrace for A {
    fn with_trace(self, t: TraceBox) -> Traced<Self> {
        Traced { a: self, t }
    }
}

impl<A> Traced<A> {
    fn wrapped(self, event: &str) -> Self {
        Traced {
            a: self.a,
            t: self.t.wrapped(event),
        }
    }

    pub fn after_resolution(self) -> Self {
        self.after(TraceBox::resolution())
    }

    fn after(self, t: TraceBox) -> Self {
        Traced {
            a: self.a,
            t: self.t.after(t),
        }
    }

    fn before(self, t: TraceBox) -> Self {
        Traced {
            a: self.a,
            t: self.t.before(t),
        }
    }

    pub fn length(&self) -> usize {
        self.t.length()
    }

    pub fn width(&self) -> usize {
        self.t.width()
    }
}

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

impl Functor for TracedClass {
    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,
    {
        f(fa.a).with_trace(fa.t)
    }

    fn replace<'a, A: 'a, B: 'a>(fa: Self::F<'a, A>, b: B) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        b.with_trace(fa.t)
    }

    fn void<'a, A: 'a>(fa: Self::F<'a, A>) -> Self::F<'a, ()>
    where
        Self: 'a,
    {
        ().with_trace(fa.t)
    }
}

impl Pure for TracedClass {
    fn pure<'a, A: 'a>(a: A) -> Self::F<'a, A> {
        a.with_trace(TraceBox::pure())
    }
}

impl ApplicativeSeq for TracedClass {
    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.a)(fa.a).with_trace(TraceBox::parallel(ff.t, fa.t))
    }
}

impl ApplicativeLA2 for TracedClass {
    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,
    {
        f(fa.a, fb.a).with_trace(TraceBox::parallel(fa.t, fb.t))
    }
}

impl ApplicativeTuple for TracedClass {
    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,
    {
        (fa.a, fb.a).with_trace(TraceBox::parallel(fa.t, fb.t))
    }
}

impl Applicative for TracedClass {
    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,
    {
        fb.a.with_trace(TraceBox::parallel(fa.t, fb.t))
    }

    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,
    {
        fa.a.with_trace(TraceBox::parallel(fa.t, fb.t))
    }
}

impl Monad for TracedClass {
    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,
    {
        f(fa.a).after(fa.t)
    }

    fn ibind<'a, A: 'a, B: 'a>(
        mut a: A,
        mut f: impl 'a + FnMut(A) -> Self::F<'a, IState<A, B>>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        let mut t = TraceBox::pure();
        loop {
            let fa = f(a);
            t = fa.t.after(t);
            match fa.a {
                IState::Pending(next_a) => a = next_a,
                IState::Done(b) => return b.with_trace(t),
            }
        }
    }

    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,
    {
        ffa.a.after(ffa.t)
    }
}

impl Diagnostic<TracedClass> for TracedDiagnostic {
    fn after<A>(fa: Traced<A>, event: &str) -> Traced<A> {
        fa.after(TraceBox::event(event))
    }

    fn before<A>(fa: Traced<A>, event: &str) -> Traced<A> {
        fa.before(TraceBox::event(event))
    }

    fn wrapped<A>(fa: Traced<A>, event: &str) -> Traced<A> {
        fa.wrapped(event)
    }
}