mod tree_node;
mod util;

use crate::util::MaybeDangling;
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::sync::Arc;

use event_listener::EventListener;
use pin_project_lite::pin_project;

/// A token which can be used to signal a cancellation request to one or more
/// tasks.
///
/// Tasks can call [`CancelToken::cancelled()`] in order to
/// obtain a Future which will be resolved when cancellation is requested.
///
/// Cancellation can be requested through the [`CancelToken::cancel`] method.
pub struct CancelToken {
    inner: Arc<tree_node::TreeNode>,
}

impl std::panic::UnwindSafe for CancelToken {}
impl std::panic::RefUnwindSafe for CancelToken {}

pin_project! {
    /// A Future that is resolved once the corresponding [`CancelToken`]
    /// is cancelled.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFuture<'a> {
        cancellation_token: &'a CancelToken,
        #[pin]
        future: EventListener<()>,
    }
}

pin_project! {
    /// A Future that is resolved once the corresponding [`CancelToken`]
    /// is cancelled.
    ///
    /// This is the counterpart to [`WaitForCancellationFuture`] that takes
    /// [`CancelToken`] by value instead of using a reference.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFutureOwned {
        // This field internally has a reference to the cancellation token, but camouflages
        // the relationship with `'static`. To avoid Undefined Behavior, we must ensure
        // that the reference is only used while the cancellation token is still alive. To
        // do that, we ensure that the future is the first field, so that it is dropped
        // before the cancellation token.
        //
        // We use `MaybeDanglingFuture` here because without it, the compiler could assert
        // the reference inside `future` to be valid even after the destructor of that
        // field runs. (Specifically, when the `WaitForCancellationFutureOwned` is passed
        // as an argument to a function, the reference can be asserted to be valid for the
        // rest of that function.) To avoid that, we use `MaybeDangling` which tells the
        // compiler that the reference stored inside it might not be valid.
        //
        // See <https://users.rust-lang.org/t/unsafe-code-review-semi-owning-weak-rwlock-t-guard/95706>
        // for more info.
        #[pin]
        future: MaybeDangling<EventListener<()>>,
        cancellation_token: CancelToken,
    }
}

impl core::fmt::Debug for CancelToken {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("CancelToken")
            .field("cancelled", &self.is_cancelled())
            .finish()
    }
}

impl Clone for CancelToken {
    /// Creates a clone of the `CancelToken` which will get cancelled
    /// whenever the current token gets cancelled, and vice versa.
    fn clone(&self) -> Self {
        self.inner.increase_handle_refcount();
        CancelToken {
            inner: self.inner.clone(),
        }
    }
}

impl Drop for CancelToken {
    fn drop(&mut self) {
        self.inner.decrease_handle_refcount();
    }
}

impl Default for CancelToken {
    fn default() -> CancelToken {
        CancelToken::new()
    }
}

impl CancelToken {
    /// Creates a new `CancelToken` in the non-cancelled state.
    pub fn new() -> CancelToken {
        CancelToken {
            inner: Arc::new(tree_node::TreeNode::new()),
        }
    }

    /// Creates a `CancelToken` which will get cancelled whenever the
    /// current token gets cancelled. Unlike a cloned `CancelToken`,
    /// cancelling a child token does not cancel the parent token.
    ///
    /// If the current token is already cancelled, the child token will get
    /// returned in cancelled state.
    pub fn child_token(&self) -> CancelToken {
        CancelToken {
            inner: self.inner.child_node(),
        }
    }

    /// Cancel the [`CancelToken`] and all child tokens which had been
    /// derived from it.
    ///
    /// This will wake up all tasks which are waiting for cancellation.
    ///
    /// Be aware that cancellation is not an atomic operation. It is possible
    /// for another thread running in parallel with a call to `cancel` to first
    /// receive `true` from `is_cancelled` on one child node, and then receive
    /// `false` from `is_cancelled` on another child node. However, once the
    /// call to `cancel` returns, all child nodes have been fully cancelled.
    pub fn cancel(&self) {
        self.inner.cancel();
    }

    /// Returns `true` if the `CancelToken` is cancelled.
    pub fn is_cancelled(&self) -> bool {
        self.inner.is_cancelled()
    }

    /// Returns a `Future` that gets fulfilled when cancellation is requested.
    ///
    /// The future will complete immediately if the token is already cancelled
    /// when this method is called.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe.
    pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
        WaitForCancellationFuture {
            cancellation_token: self,
            future: self.inner.notified(),
        }
    }

    /// Returns a `Future` that gets fulfilled when cancellation is requested.
    ///
    /// The future will complete immediately if the token is already cancelled
    /// when this method is called.
    ///
    /// The function takes self by value and returns a future that owns the
    /// token.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe.
    pub fn cancelled_owned(self) -> WaitForCancellationFutureOwned {
        WaitForCancellationFutureOwned::new(self)
    }

    /// Creates a `DropGuard` for this token.
    ///
    /// Returned guard will cancel this token (and all its children) on drop
    /// unless disarmed.
    pub fn drop_guard(self) -> DropGuard {
        DropGuard::new(self)
    }

    /// Runs a future to completion and returns its result wrapped inside of an `Option`
    /// unless the `CancelToken` is cancelled. In that case the function returns
    /// `None` and the future gets dropped.
    ///
    /// # Cancel safety
    ///
    /// This method is only cancel safe if `fut` is cancel safe.
    pub async fn run_until_cancelled<F>(&self, fut: F) -> Option<F::Output>
    where
        F: Future,
    {
        pin_project! {
            /// A Future that is resolved once the corresponding [`CancelToken`]
            /// is cancelled or a given Future gets resolved. It is biased towards the
            /// Future completion.
            #[must_use = "futures do nothing unless polled"]
            struct RunUntilCancelledFuture<'a, F: Future> {
                #[pin]
                cancellation: WaitForCancellationFuture<'a>,
                #[pin]
                future: F,
            }
        }

        impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> {
            type Output = Option<F::Output>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let this = self.project();
                if let Poll::Ready(res) = this.future.poll(cx) {
                    Poll::Ready(Some(res))
                } else if this.cancellation.poll(cx).is_ready() {
                    Poll::Ready(None)
                } else {
                    Poll::Pending
                }
            }
        }

        RunUntilCancelledFuture {
            cancellation: self.cancelled(),
            future: fut,
        }
        .await
    }

    /// Runs a future to completion and returns its result wrapped inside of an `Option`
    /// unless the `CancelToken` is cancelled. In that case the function returns
    /// `None` and the future gets dropped.
    ///
    /// The function takes self by value and returns a future that owns the token.
    ///
    /// # Cancel safety
    ///
    /// This method is only cancel safe if `fut` is cancel safe.
    pub async fn run_until_cancelled_owned<F>(self, fut: F) -> Option<F::Output>
    where
        F: Future,
    {
        self.run_until_cancelled(fut).await
    }
}

impl<'a> core::fmt::Debug for WaitForCancellationFuture<'a> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("WaitForCancellationFuture")
            .finish_non_exhaustive()
    }
}

impl<'a> Future for WaitForCancellationFuture<'a> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let mut this = self.project();
        loop {
            if this.cancellation_token.is_cancelled() {
                return Poll::Ready(());
            }

            // No wakeups can be lost here because there is always a call to
            // `is_cancelled` between the creation of the future and the call to
            // `poll`, and the code that sets the cancelled flag does so before
            // waking the `Notified`.
            if this.future.as_mut().poll(cx).is_pending() {
                return Poll::Pending;
            }

            this.future.set(this.cancellation_token.inner.notified());
        }
    }
}

impl core::fmt::Debug for WaitForCancellationFutureOwned {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("WaitForCancellationFutureOwned")
            .finish_non_exhaustive()
    }
}

impl WaitForCancellationFutureOwned {
    fn new(cancellation_token: CancelToken) -> Self {
        WaitForCancellationFutureOwned {
            // cancellation_token holds a heap allocation and is guaranteed to have a
            // stable deref, thus it would be ok to move the cancellation_token while
            // the future holds a reference to it.

            // Safety: Cancellation_token is dropped after future due to the field ordering.
            future: MaybeDangling::new(unsafe { Self::new_future(&cancellation_token) }),
            cancellation_token,
        }
    }

    /// # Safety
    /// The returned future must be destroyed before the cancellation token is
    /// destroyed.
    unsafe fn new_future(cancellation_token: &CancelToken) -> EventListener<()> {
        let inner_ptr = Arc::as_ptr(&cancellation_token.inner);
        // SAFETY: The `Arc::as_ptr` method guarantees that `inner_ptr` remains
        // valid until the strong count of the Arc drops to zero, and the caller
        // guarantees that they will drop the future before that happens.
        unsafe { (*inner_ptr).notified() }
    }
}

impl Future for WaitForCancellationFutureOwned {
    type Output = ();

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

        loop {
            if this.cancellation_token.is_cancelled() {
                return Poll::Ready(());
            }

            // No wakeups can be lost here because there is always a call to
            // `is_cancelled` between the creation of the future and the call to
            // `poll`, and the code that sets the cancelled flag does so before
            // waking the `Notified`.
            if this.future.as_mut().poll(cx).is_pending() {
                return Poll::Pending;
            }

            // SAFETY: Cancellation_token is dropped after future due to the field ordering.
            this.future.set(MaybeDangling::new(unsafe {
                Self::new_future(this.cancellation_token)
            }));
        }
    }
}

pub struct DropGuard {
    token: Option<CancelToken>,
}

impl DropGuard {
    const fn new(token: CancelToken) -> Self {
        Self { token: Some(token) }
    }

    pub fn disarm(&mut self) {
        self.token.take();
    }
}

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