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

use crate::{higher_ranked_trait, higher_ranked_type};

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

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

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, 'ctx, 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, 'ctx, Self::Future<T>, T>;

    #[cfg(feature = "alloc")]
    #[inline]
    fn wrap_boxed<'a, F>(
        future: core::pin::Pin<Box<F>>,
    ) -> SendFuture::T<'a, 'ctx, 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, 'ctx, T, E> = SendFuture::T<'a, 'ctx, <E as Effect<'ctx>>::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: &'static 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) }
}

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

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

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, 'ctx, 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, 'ctx, Self::Future<T>, T> {
        core::future::ready(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")]
higher_ranked_type! {
    pub type BoxFuture['ctx, Output]: (SendFuture)[Output]
    where {
        Output: Send,
    } = for<'lt> sealed::BoxedFuture<'lt, Output>
}

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

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

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

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