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

use crate::{higher_ranked_trait, higher_ranked_type, hkt::Marker};

higher_ranked_trait! {
    pub type class SendFuture[Output] for<'a> {
        type Bound = &'a Output;

        type T: { core::future::Future<Output = Output> + Sized + Send + 'a }
        where {
            Output: 'a
        };
        type HigherRanked: {} 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, F::Output, Self::Future<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, T, Self::Future<T>>;

    fn map<'a, T, U, F>(
        future: SendFuture::T<'a, T, Self::Future<T>>,
        func: F,
    ) -> SendFuture::T<'a, U, Self::Future<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, F::Output, Self::Future<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, T, <E as Effect>::Future<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) }
}

higher_ranked_type! {
    impl SendFuture {
        impl['a, Output] type T['a, Output] for core::future::Ready<Output> =
            core::future::Ready<Output>
        where {
            Output: Send
        };

        impl['a, Output] type HigherRanked['a, Output] for core::future::Ready<Output> =
            core::future::Ready<Output>
        where {
            Output: Send
        };
    }
}

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

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, F::Output, Self::Future<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, T, Self::Future<T>> {
        core::future::ready(value)
    }

    fn map<'a, T, U, F>(
        future: SendFuture::T<'a, T, Self::Future<T>>,
        func: F,
    ) -> SendFuture::T<'a, U, Self::Future<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")]
pub struct BoxedFutureHrt<Output>(Marker<Output>);

#[cfg(feature = "alloc")]
higher_ranked_type! {
    impl SendFuture {
        impl['a, Output] type T['a, Output] for BoxedFutureHrt<Output> =
            sealed::BoxedFuture<'a, Output>
        where {
            Output: Send
        };

        impl['a, Output] type HigherRanked['a, Output] for sealed::BoxedFuture<'a, Output> =
            BoxedFutureHrt<Output>
        where {
            Output: Send
        };
    }
}

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

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

    fn wrap<'a, F>(future: F) -> SendFuture::T<'a, F::Output, Self::Future<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, T, Self::Future<T>> {
        sealed::BoxedFuture::Ready(core::future::ready(value))
    }

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