use std::cmp::max;

use super::*;
use crate::core::*;
use crate::func::*;

pub struct TracedDiagnostic;

pub struct TestContextTraced;

impl Context for TestContextTraced {
    type T = TracedClass;

    type D = TracedDiagnostic;

    type LookupError<'a> = TestLookupError<'a>;

    fn hash(s: &[u8]) -> Hash {
        TestContextPlain::hash(s)
    }
}

pub struct TracedClass;

pub enum Trace {
    Pure,
    InvolvesOneResolution,
    Event(String),
    Parallel(Box<Trace>, Box<Trace>),
    Sequential {
        first: Box<Trace>,
        second: Box<Trace>,
    },
    Wrapped {
        name: String,
        trace: Box<Trace>,
    },
}

pub struct Traced<A> {
    a: A,
    t: Box<Trace>,
}

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

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

impl Trace {
    fn pure() -> Box<Self> {
        Self::Pure.into()
    }

    fn resolution() -> Box<Self> {
        Self::InvolvesOneResolution.into()
    }

    fn event(event: &str) -> Box<Self> {
        Self::Event(event.into()).into()
    }

    fn wrapped(self: Box<Self>, event: &str) -> Box<Self> {
        Self::Wrapped {
            name: event.into(),
            trace: self,
        }
        .into()
    }

    fn after(self: Box<Self>, t: Box<Self>) -> Box<Self> {
        match (*self, *t) {
            (Self::Pure, a) => a.into(),
            (a, Self::Pure) => a.into(),
            (a, b) => Self::Sequential {
                first: b.into(),
                second: a.into(),
            }
            .into(),
        }
    }

    fn before(self: Box<Self>, t: Box<Self>) -> Box<Self> {
        match (*self, *t) {
            (Self::Pure, a) => a.into(),
            (a, Self::Pure) => a.into(),
            (a, b) => Self::Sequential {
                first: a.into(),
                second: b.into(),
            }
            .into(),
        }
    }

    fn parallel(ta: Box<Self>, tb: Box<Self>) -> Box<Self> {
        match (*ta, *tb) {
            (Self::Pure, a) => a.into(),
            (a, Self::Pure) => a.into(),
            (a, b) => Self::Parallel(a.into(), b.into()).into(),
        }
    }

    fn length(&self) -> usize {
        match self {
            Self::Pure => 0,
            Self::InvolvesOneResolution => 1,
            Self::Event(_) => 0,
            Self::Parallel(a, b) => max(a.length(), b.length()),
            Self::Sequential { first, second } => first.length() + second.length(),
            Self::Wrapped { name: _, trace } => trace.length(),
        }
    }

    fn width(&self) -> usize {
        match self {
            Self::Pure => 0,
            Self::InvolvesOneResolution => 1,
            Self::Event(_) => 0,
            Self::Parallel(a, b) => a.width() + b.width(),
            Self::Sequential { first, second } => max(first.width(), second.width()),
            Self::Wrapped { name: _, trace } => trace.width(),
        }
    }
}

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

    fn before(self, t: Box<Trace>) -> 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 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(Trace::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(Trace::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(Trace::parallel(fa.t, fb.t))
    }
}

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

    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(Trace::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(Trace::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 = Trace::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)
    }
}

struct TracedResolver<'a> {
    resolver: Rc<dyn Resolver<'a, TestContextTraced>>,
}

impl<'a> TracedResolver<'a> {
    fn new(
        resolver: Rc<dyn Resolver<'a, TestContextTraced>>,
    ) -> Rc<dyn Resolver<'a, TestContextTraced>> {
        Rc::new(Self { resolver })
    }
}

impl<'a> Resolver<'a, TestContextTraced> for TracedResolver<'a> {
    fn resolve(self: Rc<Self>, address: Address) -> HashResolution<'a, TestContextTraced> {
        TracedClass::fmap(
            |resolved| {
                let (src, resolver) = resolved?;
                let delayed: Rc<dyn Resolver<'a, TestContextTraced>> =
                    Rc::new(TracedResolver { resolver });
                Ok((src, delayed))
            },
            self.resolver.clone().resolve(address),
        )
        .after(Trace::resolution())
    }
}

pub trait Traceable<'a>: Mentionable<'a, TestContextTraced> + Sized {
    fn trace(self: Rc<Self>) -> CastResult<'a, TestContextTraced, Self>;
}

impl<'a, A: Mentionable<'a, TestContextTraced>> Traceable<'a> for A {
    fn trace(self: Rc<Self>) -> CastResult<'a, TestContextTraced, Self> {
        let factory = self.factory();
        TypelessMentionable::from_typed(self).cast_full(factory, TracedResolver::new)
    }
}

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

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

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