dope-runtime 0.2.3

use std::{
    future::Future,
    io,
    pin::Pin,
    ptr::NonNull,
    task::{Context, Poll},
    time::{Duration, Instant},
};

use dope_core::driver::CompletionBackend;

use super::executor;

pub async fn sleep_with<B: CompletionBackend>(
    ctx: executor::RuntimeCtx<B>,
    duration: Duration,
) -> io::Result<()> {
    if duration.is_zero() {
        return Ok(());
    }
    Sleep::new(ctx.timers(), Instant::now() + duration).await;
    Ok(())
}

pub async fn sleep<B: CompletionBackend>(duration: Duration) -> io::Result<()> {
    sleep_with(executor::current::<B>(), duration).await
}

pub fn timeout_with<B, F>(
    ctx: executor::RuntimeCtx<B>,
    duration: Duration,
    fut: F,
) -> impl Future<Output = io::Result<F::Output>>
where
    B: CompletionBackend,
    F: Future,
{
    Timeout::new(ctx.timers(), duration, fut)
}

pub fn timeout<B, F>(duration: Duration, fut: F) -> impl Future<Output = io::Result<F::Output>>
where
    B: CompletionBackend,
    F: Future,
{
    timeout_with(executor::current::<B>(), duration, fut)
}

struct Timeout<F>
where
    F: Future,
{
    fut: F,
    deadline: Instant,
    sleep: Sleep,
}

impl<F> Future for Timeout<F>
where
    F: Future,
{
    type Output = io::Result<F::Output>;

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

        if let Poll::Ready(v) = unsafe { Pin::new_unchecked(&mut this.fut) }.poll(cx) {
            return Poll::Ready(Ok(v));
        }

        let now = Instant::now();
        if now >= this.deadline {
            return Poll::Ready(Err(io::Error::new(io::ErrorKind::TimedOut, "timeout")));
        }

        match Pin::new(&mut this.sleep).poll(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(()) => Poll::Ready(Err(io::Error::new(io::ErrorKind::TimedOut, "timeout"))),
        }
    }
}

impl<F> Timeout<F>
where
    F: Future,
{
    fn new(timers: NonNull<executor::TimerHeap>, duration: Duration, fut: F) -> Self {
        let deadline = Instant::now() + duration;
        Self {
            fut,
            deadline,
            sleep: Sleep::new(timers, deadline),
        }
    }
}

struct Sleep {
    deadline: Instant,
    id: Option<u64>,
    timers: NonNull<executor::TimerHeap>,
}

impl Sleep {
    fn new(timers: NonNull<executor::TimerHeap>, deadline: Instant) -> Self {
        Self {
            deadline,
            id: None,
            timers,
        }
    }

    #[inline(always)]
    fn timers(&self) -> &executor::TimerHeap {
        unsafe { self.timers.as_ref() }
    }
}

impl Drop for Sleep {
    fn drop(&mut self) {
        if let Some(id) = self.id.take() {
            self.timers().cancel(id);
        }
    }
}

impl Future for Sleep {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let now = Instant::now();
        if now >= self.deadline {
            if let Some(id) = self.id.take() {
                self.timers().cancel(id);
            }
            return Poll::Ready(());
        }

        let deadline = self.deadline;
        let id = self.id;
        let new_id = self.timers().register(id, deadline, cx.waker());
        self.id = Some(new_id);
        Poll::Pending
    }
}