#![forbid(unsafe_code)]

use crate::schedule_wake;
use core::fmt::{Debug, Display, Formatter};
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};
use std::error::Error;
use std::sync::{Arc, Mutex};
use std::time::Instant;

/// - `DeadlineError::TimerThreadNotStarted`
/// - `DeadlineError::DeadlineExceeded`
#[derive(Debug, PartialEq)]
pub enum DeadlineError {
    TimerThreadNotStarted,
    DeadlineExceeded,
}
impl From<DeadlineError> for std::io::Error {
    fn from(error: DeadlineError) -> Self {
        match error {
            DeadlineError::TimerThreadNotStarted => {
                std::io::Error::new(std::io::ErrorKind::Other, "TimerThreadNotStarted")
            }
            DeadlineError::DeadlineExceeded => {
                std::io::Error::new(std::io::ErrorKind::TimedOut, "DeadlineExceeded")
            }
        }
    }
}
impl Display for DeadlineError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        std::fmt::Debug::fmt(self, f)
    }
}
impl Error for DeadlineError {}

#[derive(Debug, PartialEq)]
pub struct DeadlineExceeded;
impl From<DeadlineExceeded> for std::io::Error {
    fn from(_error: DeadlineExceeded) -> Self {
        std::io::Error::new(std::io::ErrorKind::TimedOut, "DeadlineExceeded")
    }
}
impl Display for DeadlineExceeded {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        std::fmt::Debug::fmt(self, f)
    }
}
impl Error for DeadlineExceeded {}

/// A future wrapper that returns DeadlineExceeded at a specified deadline.
///
/// It is returned by [`with_deadline`] and [`with_timeout`].
#[must_use = "futures stay idle unless you await them"]
pub struct DeadlineFuture<R, Fut: Future<Output = R> + Send + Unpin + 'static> {
    inner: Fut,
    deadline: std::time::Instant,
    waker: Arc<Mutex<Option<Waker>>>,
}
impl<R, Fut: Future<Output = R> + Send + Unpin + 'static> DeadlineFuture<R, Fut> {
    /// Makes a future that awaits `inner`,
    /// but returns [`DeadlineError`](enum.DeadlineError.html) after `deadline`.
    ///
    /// Note that `inner` must be
    /// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
    /// Use [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
    /// to make it Unpin.
    /// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
    /// to do it with unsafe code that does not allocate memory.
    pub fn new(inner: Fut, deadline: Instant) -> Self {
        Self {
            inner,
            deadline,
            waker: Arc::new(Mutex::new(None)),
        }
    }
}
impl<R, Fut: Future<Output = R> + Send + Unpin + 'static> Future for DeadlineFuture<R, Fut> {
    type Output = Result<R, DeadlineError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut_self = self.get_mut();
        // The primary purpose of deadlines is to shed load during overload.
        // If the inner future completed and the deadline exceeded, the process
        // is likely overloaded.  In this case, we return error to shed load.
        if mut_self.deadline < std::time::Instant::now() {
            return Poll::Ready(Err(DeadlineError::DeadlineExceeded));
        } else {
            match Pin::new(&mut mut_self.inner).poll(cx) {
                Poll::Ready(r) => return Poll::Ready(Ok(r)),
                Poll::Pending => {}
            }
        }
        let old_waker = mut_self.waker.lock().unwrap().replace(cx.waker().clone());
        if old_waker.is_none() {
            schedule_wake(mut_self.deadline, mut_self.waker.clone())
                .map_err(|_| DeadlineError::TimerThreadNotStarted)?;
        }
        Poll::Pending
    }
}

/// Awaits `inner`, but returns [`DeadlineExceeded`](struct.DeadlineExceeded.html)
/// after `deadline`.
///
/// First moves `inner` to the heap, to make it
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// Use
/// [`DeadlineFuture::new`](https://docs.rs/safina-timer/latest/safina_timer/struct.DeadlineFuture.html)
/// to avoid allocating on the heap.
///
/// Panics if [`start_timer_thread()`](fn.start_timer_thread.html) has not been called.
/// If you need to handle this error, use
/// [`DeadlineFuture::new`](https://docs.rs/safina-timer/latest/safina_timer/struct.DeadlineFuture.html).
pub async fn with_deadline<R, Fut: Future<Output = R> + Send + 'static>(
    inner: Fut,
    deadline: std::time::Instant,
) -> Result<R, DeadlineExceeded> {
    match DeadlineFuture::new(Box::pin(inner), deadline).await {
        Ok(result) => Ok(result),
        Err(DeadlineError::DeadlineExceeded) => Err(DeadlineExceeded),
        Err(DeadlineError::TimerThreadNotStarted) => panic!("TimerThreadNotStarted"),
    }
}

/// Awaits `inner`, but returns [`DeadlineExceeded`](struct.DeadlineExceeded.html)
/// after `duration` time from now.
///
/// First moves `inner` to the heap, to make it
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// Use
/// [`DeadlineFuture::new`](https://docs.rs/safina-timer/latest/safina_timer/struct.DeadlineFuture.html)
/// to avoid allocating on the heap.
///
/// Panics if [`start_timer_thread()`](fn.start_timer_thread.html) has not been called.
/// If you need to handle this error, use
/// [`DeadlineFuture::new`](https://docs.rs/safina-timer/latest/safina_timer/struct.DeadlineFuture.html).
pub async fn with_timeout<R, Fut: Future<Output = R> + Send + 'static>(
    inner: Fut,
    duration: std::time::Duration,
) -> Result<R, DeadlineExceeded> {
    with_deadline(inner, Instant::now() + duration).await
}

#[cfg(test)]
mod tests {
    use super::super::*;
    use core::future::Future;
    use core::pin::Pin;
    use core::task::{Context, Poll};
    use core::time::Duration;
    use rusty_fork::rusty_fork_test;
    use safina_async_test::async_test;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::time::Instant;

    struct PendingFuture;
    impl Future for PendingFuture {
        type Output = ();
        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            Poll::Pending
        }
    }

    fn timer_thread_not_started_inner() {
        let deadline = Instant::now() + Duration::from_millis(1000);
        assert_eq!(
            "TimerThreadNotStarted",
            *(std::panic::catch_unwind(|| {
                safina_executor::block_on(async move {
                    with_deadline(async { PendingFuture {}.await }, deadline).await
                })
            })
            .unwrap_err()
            .downcast::<&'static str>()
            .unwrap())
        );

        for _ in 0..2 {
            start_timer_thread();
            assert_eq!(
                DeadlineExceeded,
                safina_executor::block_on(async move {
                    with_deadline(async { PendingFuture {}.await }, deadline).await
                })
                .unwrap_err()
            );
        }
    }
    rusty_fork_test! {
        #[test]
        fn timer_thread_not_started() {
            timer_thread_not_started_inner();
        }
    }

    #[async_test]
    async fn with_deadline_should_timeout() {
        start_timer_thread();
        let before = Instant::now();
        assert_eq!(
            DeadlineExceeded,
            with_deadline(
                async { sleep_for(Duration::from_millis(200)).await },
                before + Duration::from_millis(100)
            )
            .await
            .unwrap_err()
        );
        expect_elapsed(before, 100..200);
    }

    #[async_test]
    async fn with_deadline_should_return_result() {
        start_timer_thread();
        let before = Instant::now();
        assert_eq!(
            42u8,
            with_deadline(
                async {
                    sleep_for(Duration::from_millis(100)).await;
                    42u8
                },
                before + Duration::from_millis(200)
            )
            .await
            .unwrap()
        );
        expect_elapsed(before, 100..200);
    }

    #[async_test]
    async fn with_timeout_should_timeout() {
        start_timer_thread();
        let before = Instant::now();
        assert_eq!(
            DeadlineExceeded,
            with_timeout(
                async { sleep_for(Duration::from_millis(200)).await },
                Duration::from_millis(100)
            )
            .await
            .unwrap_err()
        );
        expect_elapsed(before, 100..200);
    }

    #[async_test]
    async fn with_timeout_should_return_result() {
        start_timer_thread();
        let before = Instant::now();
        assert_eq!(
            42u8,
            with_timeout(
                async {
                    sleep_for(Duration::from_millis(100)).await;
                    42u8
                },
                Duration::from_millis(200)
            )
            .await
            .unwrap()
        );
        expect_elapsed(before, 100..200);
    }

    #[async_test]
    async fn should_return_immediately_when_inner_already_ready() {
        start_timer_thread();
        let before = Instant::now();
        with_deadline(async { 42u8 }, before + Duration::from_millis(100))
            .await
            .unwrap();
        expect_elapsed(before, 0..20);
    }

    #[async_test]
    async fn deadline_already_past() {
        start_timer_thread();
        let before = Instant::now();
        assert_eq!(
            DeadlineExceeded,
            with_deadline(
                async { sleep_for(Duration::from_millis(200)).await },
                before - Duration::from_millis(100)
            )
            .await
            .unwrap_err()
        );
        expect_elapsed(before, 0..20);
    }

    #[test]
    pub fn should_use_most_recent_waker_passed_to_poll() {
        // "Note that on multiple calls to poll, only the Waker from the Context
        // passed to the most recent call should be scheduled to receive a wakeup."
        // https://doc.rust-lang.org/stable/std/future/trait.Future.html#tymethod.poll
        start_timer_thread();
        let deadline = Instant::now() + Duration::from_millis(100);
        let mut fut =
            Box::pin(
                async move { with_deadline(async { PendingFuture {}.await }, deadline).await },
            );
        let waker1_called = Arc::new(AtomicBool::new(false));
        {
            let waker1 = FakeWaker::new(&waker1_called).into_waker();
            let mut cx = Context::from_waker(&waker1);
            assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
        }
        let waker2_called = Arc::new(AtomicBool::new(false));
        {
            let waker2 = FakeWaker::new(&waker2_called).into_waker();
            let mut cx = Context::from_waker(&waker2);
            assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
        }
        std::thread::sleep(Duration::from_millis(200));
        {
            let waker3_called = Arc::new(AtomicBool::new(true /* should never get called */));
            let waker3 = FakeWaker::new(&waker3_called).into_waker();
            let mut cx = Context::from_waker(&waker3);
            assert_eq!(
                Poll::Ready(Err(DeadlineExceeded)),
                fut.as_mut().poll(&mut cx)
            );
        }
        assert!(!waker1_called.load(Ordering::Acquire));
        assert!(waker2_called.load(Ordering::Acquire));
    }
}