use core::{
    marker::PhantomData,
    pin::pin,
    ptr,
    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};

use crate::{bijective_higher_ranked_trait, bijective_higher_ranked_type};

bijective_higher_ranked_trait! {
    pub type class SendFuture[][Output]: [for<'lt> core::future::Future<Output = Output> + Sized + Send + 'lt]
    where {
        Output: Send,
    }
}

/// Trait for effects.
pub trait Effect: Send + 'static {
    type Future<T: Send>: SendFuture::MemberType<T>;

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, Self::Future<F::Output>, F::Output>
    where
        F: core::future::Future + Send + 'a,
        <F as core::future::Future>::Output: Send;

    fn ready<'a, T: Send>(value: T) -> SendFuture::T<'a, Self::Future<T>, T>;

    fn map<'a, T, U, F>(
        future: SendFuture::T<'a, Self::Future<T>, T>,
        func: F,
    ) -> SendFuture::T<'a, Self::Future<U>, U>
    where
        T: Send,
        U: Send,
        F: FnOnce(T) -> U + Send + 'a;

    #[cfg(feature = "alloc")]
    #[inline]
    fn wrap_boxed<'a, F>(
        future: core::pin::Pin<Box<F>>,
    ) -> SendFuture::T<'a, Self::Future<F::Output>, F::Output>
    where
        F: core::future::Future + Send + 'a,
        <F as core::future::Future>::Output: Send,
    {
        Self::wrap(future)
    }
}

pub type Future<'a, T, E> = SendFuture::T<'a, <E as Effect>::Future<T>, T>;

pub struct Blocking<B = Spin> {
    _marker: PhantomData<fn() -> B>,
}

pub trait BlockOn: 'static {
    fn block_on<F>(future: F) -> F::Output
    where
        F: core::future::Future + Send,
        <F as core::future::Future>::Output: Send;
}

/// [`BlockOn`] implementer that just spins on the future.
///
/// This is useful for futures that are alwayd ready.
pub enum Spin {}

impl BlockOn for Spin {
    #[inline(always)]
    fn block_on<F>(future: F) -> F::Output
    where
        F: core::future::Future + Send,
    {
        let waker = noop();
        let mut context = Context::from_waker(&waker);

        let mut future = pin!(future);
        loop {
            if let Poll::Ready(value) = future.as_mut().poll(&mut context) {
                return value;
            }
        }
    }
}

#[inline]
pub fn noop() -> Waker {
    const VTABLE: &RawWakerVTable = &RawWakerVTable::new(
        // Cloning just returns a new no-op raw waker
        |_| RAW,
        // `wake` does nothing
        |_| {},
        // `wake_by_ref` does nothing
        |_| {},
        // Dropping does nothing as we don't allocate anything
        |_| {},
    );
    const RAW: RawWaker = RawWaker::new(ptr::null(), VTABLE);
    unsafe { Waker::from_raw(RAW) }
}

bijective_higher_ranked_type! {
    pub type ReadyFuture[][Output]: SendFuture[][Output]
    for<'lt> (core::future::Ready<Output>)
    where {
        Output: Send,
    }
}

impl<B: BlockOn> Effect for Blocking<B> {
    type Future<T: Send> = ReadyFuture<T>;

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, Self::Future<F::Output>, F::Output>
    where
        F: core::future::Future + Send + 'a,
        <F as core::future::Future>::Output: Send,
    {
        core::future::ready(B::block_on(future))
    }

    fn ready<'a, T: Send>(value: T) -> SendFuture::T<'a, Self::Future<T>, T> {
        core::future::ready(value)
    }

    fn map<'a, T, U, F>(
        future: SendFuture::T<'a, Self::Future<T>, T>,
        func: F,
    ) -> SendFuture::T<'a, Self::Future<U>, U>
    where
        T: Send,
        U: Send,
        F: FnOnce(T) -> U + Send + 'a,
    {
        let value = B::block_on(future);
        core::future::ready(func(value))
    }
}

mod sealed {
    pub enum BoxedFuture<'lt, Output> {
        Box(core::pin::Pin<Box<dyn core::future::Future<Output = Output> + Send + 'lt>>),
        Ready(core::future::Ready<Output>),
    }

    impl<'lt, Output> core::future::Future for BoxedFuture<'lt, Output> {
        type Output = Output;

        fn poll(
            mut self: core::pin::Pin<&mut Self>,
            cx: &mut core::task::Context<'_>,
        ) -> core::task::Poll<Self::Output> {
            match &mut *self {
                BoxedFuture::Box(future) => future.as_mut().poll(cx),
                BoxedFuture::Ready(future) => core::pin::Pin::new(future).poll(cx),
            }
        }
    }
}

#[cfg(feature = "alloc")]
bijective_higher_ranked_type! {
    pub type BoxFuture[][Output]: SendFuture[][Output]
    for<'lt> (sealed::BoxedFuture<'lt, Output>)
    where {
        Output: Send,
    }
}

#[cfg(feature = "alloc")]
pub enum Async {}

#[cfg(feature = "alloc")]
impl Effect for Async {
    type Future<T: Send> = BoxFuture<T>;

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, Self::Future<F::Output>, F::Output>
    where
        F: core::future::Future + Send + 'a,
        <F as core::future::Future>::Output: Send,
    {
        sealed::BoxedFuture::Box(Box::pin(future))
    }

    fn ready<'a, T: Send>(value: T) -> SendFuture::T<'a, Self::Future<T>, T> {
        sealed::BoxedFuture::Ready(core::future::ready(value))
    }

    fn map<'a, T, U, F>(
        future: SendFuture::T<'a, Self::Future<T>, T>,
        func: F,
    ) -> SendFuture::T<'a, Self::Future<U>, U>
    where
        T: Send,
        U: Send,
        F: FnOnce(T) -> U + Send + 'a,
    {
        match future {
            sealed::BoxedFuture::Box(future) => Self::wrap(async { func(future.await) }),
            sealed::BoxedFuture::Ready(future) => {
                let value = Spin::block_on(future);
                Self::ready(func(value))
            }
        }
    }

    fn wrap_boxed<'a, F>(
        future: core::pin::Pin<Box<F>>,
    ) -> SendFuture::T<'a, Self::Future<F::Output>, F::Output>
    where
        F: core::future::Future + Send + 'a,
        <F as core::future::Future>::Output: Send,
    {
        sealed::BoxedFuture::Box(future)
    }
}