use std::cmp::max;

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

use super::*;

pub struct TestContextCounted;

impl Context for TestContextCounted {
    type T = CountedClass;

    type Fallible = classes::result::ResultFailOver<Self::T>;

    type D = NoDiagnostic;

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

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

#[derive(CovariantFunctor)]
pub struct CountedClass;

pub struct Counted<A> {
    a: A,
    n: usize,
}

trait WithCount: Sized {
    fn with_count(self, n: usize) -> Counted<Self>;
}

impl<A> WithCount for A {
    fn with_count(self, n: usize) -> Counted<Self> {
        Counted { a: self, n }
    }
}

impl<A> Counted<A> {
    fn add(self, n: usize) -> Self {
        Counted {
            a: self.a,
            n: self.n + n,
        }
    }

    pub fn count(&self) -> usize {
        self.n
    }
}

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

impl Functor for CountedClass {
    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_count(fa.n)
    }

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

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

impl Pure for CountedClass {
    fn pure<'a, A: 'a>(a: A) -> Self::F<'a, A> {
        a.with_count(0)
    }
}

impl ApplicativeSeq for CountedClass {
    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_count(max(ff.n, fa.n))
    }
}

impl ApplicativeLA2 for CountedClass {
    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_count(max(fa.n, fb.n))
    }
}

impl ApplicativeTuple for CountedClass {
    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_count(max(fa.n, fb.n))
    }
}

impl Applicative for CountedClass {
    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_count(max(fa.n, fb.n))
    }

    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_count(max(fa.n, fb.n))
    }
}

impl Monad for CountedClass {
    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).add(fa.n)
    }

    fn iterate_mut<'a, A: 'a, B: 'a>(
        mut a: A,
        mut f: impl 'a + FnMut(A) -> Self::F<'a, ControlFlow<B, A>>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        let mut n = 0;
        loop {
            let fa = f(a);
            n += fa.n;
            match fa.a {
                ControlFlow::Continue(next_a) => a = next_a,
                ControlFlow::Break(b) => return b.with_count(n),
            }
        }
    }

    fn iterate_argument<'a, A: 'a, B: 'a>(
        mut a: A,
        mut f: impl AIterative<'a, T = Self, A = A, B = B>,
    ) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        let mut n = 0;
        loop {
            let fa = f.next(a);
            n += fa.n;
            match fa.a {
                ControlFlow::Continue((next_a, next_f)) => (a, f) = (next_a, next_f),
                ControlFlow::Break(b) => return b.with_count(n),
            }
        }
    }

    fn iterate<'a, B: 'a>(mut f: impl Iterative<'a, T = Self, B = B>) -> Self::F<'a, B>
    where
        Self: 'a,
    {
        let mut n = 0;
        loop {
            let fa = f.next();
            n += fa.n;
            match fa.a {
                ControlFlow::Continue(next_f) => f = next_f,
                ControlFlow::Break(b) => return b.with_count(n),
            }
        }
    }

    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.add(ffa.n)
    }
}

struct CountedResolver<'a> {
    resolver: Rc<dyn Resolver<'a, TestContextCounted>>,
}

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

impl<'a> Resolver<'a, TestContextCounted> for CountedResolver<'a> {
    fn resolve(self: Rc<Self>, address: Address) -> HashResolution<'a, TestContextCounted> {
        CountedClass::fmap(
            |resolved| {
                let (src, resolver) = resolved?;
                let delayed: Rc<dyn Resolver<'a, TestContextCounted>> =
                    Rc::new(CountedResolver { resolver });
                Ok((src, delayed))
            },
            self.resolver.clone().resolve(address),
        )
        .add(1)
    }
}

pub trait Delayable<'a>: Mentionable<'a, TestContextCounted> + Sized {
    fn delay(self: Rc<Self>) -> CastResult<'a, TestContextCounted, Self>;
}

impl<'a, A: Mentionable<'a, TestContextCounted>> Delayable<'a> for A {
    fn delay(self: Rc<Self>) -> CastResult<'a, TestContextCounted, Self> {
        let factory = self.factory();
        TypelessMentionable::from_typed(self).cast_full(factory, CountedResolver::new)
    }
}