use std::{future::Future, pin::Pin, task};

use futures::stream;

// Combinators for futures that prevent waits, even when the underlying future
// would block.
//
// These are only useful for futures with no bound on how long they may wait,
// and this trait is only implemented on futures that are likely to have that
// characteristic. Trying to apply this to futures that already have some
// bounded wait time may make tests fail inappropriately and can hide other
// logic errors.
pub trait Expect: Sized {
    // The returned future expects the underlying future to be ready immediately,
    // and panics with the provided message if it is not.
    //
    // For stream operations, can be used to assert immediate completion.
    fn expect_ready(self, message: &str) -> Ready<'_, Self>
    where
        Self: Future;

    // The returned future expects the underlying future _not_ to be ready, and
    // panics if it is. This is usually a useful proxy for "I expect this to never
    // arrive" or "to not be here yet." The future is transformed to return `()`,
    // since the underlying future can never provide a value.
    //
    // For stream operations, can be used to assert that completion hasn't happened
    // yet.
    fn expect_wait(self, message: &str) -> Wait<'_, Self>
    where
        Self: Future;

    // The returned future expects the underlying future to resolve immediately, to
    // a `Some` value. If it resolves to `None` or is not ready, it panics. The
    // future is transformed to return the inner value from the `Some` case, like
    // [`Option::expect`].
    //
    // For stream operations, can be used to assert that the stream has at least one
    // message.
    fn expect_some<T>(self, message: &str) -> Some<'_, Self>
    where
        Self: Future<Output = Option<T>>;

    // The returned future expects the underlying future to resolve immediately, to
    // a `None` value. If it resolves to `Some(_)`, or is not ready, it panics. The
    // future is transformed to return `()`, since the underlying future's value is
    // fixed.
    //
    // For stream operations, can be used to assert that the stream has ended.
    fn expect_none<T>(self, message: &str) -> None<'_, Self>
    where
        Self: Future<Output = Option<T>>;
}

impl<St> Expect for stream::Next<'_, St> {
    fn expect_ready(self, message: &str) -> Ready<'_, Self> {
        Ready {
            future: self,
            message,
        }
    }

    fn expect_wait(self, message: &str) -> Wait<'_, Self> {
        Wait {
            future: self,
            message,
        }
    }

    fn expect_some<T>(self, message: &str) -> Some<'_, Self>
    where
        Self: Future<Output = Option<T>>,
    {
        Some {
            future: self,
            message,
        }
    }

    fn expect_none<T>(self, message: &str) -> None<'_, Self>
    where
        Self: Future<Output = Option<T>>,
    {
        None {
            future: self,
            message,
        }
    }
}

impl<St, C> Expect for stream::Collect<St, C> {
    fn expect_ready(self, message: &str) -> Ready<'_, Self> {
        Ready {
            future: self,
            message,
        }
    }

    fn expect_wait(self, message: &str) -> Wait<'_, Self> {
        Wait {
            future: self,
            message,
        }
    }

    fn expect_some<T>(self, message: &str) -> Some<'_, Self>
    where
        Self: Future<Output = Option<T>>,
    {
        Some {
            future: self,
            message,
        }
    }

    fn expect_none<T>(self, message: &str) -> None<'_, Self>
    where
        Self: Future<Output = Option<T>>,
    {
        None {
            future: self,
            message,
        }
    }
}

#[pin_project::pin_project]
pub struct Ready<'m, F> {
    #[pin]
    future: F,
    message: &'m str,
}

impl<F> Future for Ready<'_, F>
where
    F: Future + std::fmt::Debug,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();

        if let task::Poll::Ready(value) = this.future.poll(cx) {
            task::Poll::Ready(value)
        } else {
            panic!("{}", this.message);
        }
    }
}

#[pin_project::pin_project]
pub struct Wait<'m, F> {
    #[pin]
    future: F,
    message: &'m str,
}

impl<F> Future for Wait<'_, F>
where
    F: Future + std::fmt::Debug,
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();

        if this.future.poll(cx).is_pending() {
            task::Poll::Ready(())
        } else {
            panic!("{}", this.message);
        }
    }
}

#[pin_project::pin_project]
pub struct Some<'m, F> {
    #[pin]
    future: F,
    message: &'m str,
}

impl<F, T> Future for Some<'_, F>
where
    F: Future<Output = Option<T>> + std::fmt::Debug,
{
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();

        if let task::Poll::Ready(Option::Some(value)) = this.future.poll(cx) {
            task::Poll::Ready(value)
        } else {
            panic!("{}", this.message)
        }
    }
}

#[pin_project::pin_project]
pub struct None<'m, F> {
    #[pin]
    future: F,
    message: &'m str,
}

impl<F, T> Future for None<'_, F>
where
    F: Future<Output = Option<T>> + std::fmt::Debug,
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();

        if let task::Poll::Ready(Option::None) = this.future.poll(cx) {
            task::Poll::Ready(())
        } else {
            panic!("{}", this.message)
        }
    }
}