/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::{
    collections::btree_map::BTreeMap,
    fmt::Debug,
    future::Future,
    mem,
    num::NonZeroUsize,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
};

use futures::future::{self, Either};

/// An object that signals when a task, or a group of tasks,
/// should be aborted.
///
/// An [`AbortController`] works like the DOM API of the same name.
/// Each controller has one or more associated [`AbortSignal`]s, which are
/// futures that observe the controller. When the controller tells its signals
/// to abort, any tasks that are waiting on those signals can resume executing.
///
/// [`AbortController`] is similar to [`future::AbortHandle`], but can wake up
/// more than one task that's waiting on an associated [`AbortSignal`].
pub struct AbortController {
    state: Arc<Mutex<AbortState>>,
}

impl AbortController {
    /// Creates a new controller.
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(AbortState::Armed {
                next_id: NonZeroUsize::MIN,
                wakers: BTreeMap::new(),
            })),
        }
    }

    /// Creates an associated signal which can be waited on.
    pub fn signal(&self) -> AbortSignal {
        AbortSignal {
            id: AbortSignalId::Unreserved,
            state: self.state.clone(),
        }
    }

    /// Tells all associated signals to abort.
    pub fn abort(&self) {
        let wakers = match mem::replace(&mut *self.state.lock().unwrap(), AbortState::Aborted) {
            AbortState::Aborted => return,
            AbortState::Armed { wakers, .. } => wakers,
        };
        for waker in wakers.into_values() {
            waker.wake()
        }
    }
}

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

impl Debug for AbortController {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("AbortController { .. }")
    }
}

/// A future that completes when its associated [`AbortController`]
/// is aborted.
///
/// [`AbortSignal`] is similar to [`future::AbortRegistration`] and
/// [`future::Abortable`], but can be cloned and shared between
/// multiple tasks.
pub struct AbortSignal {
    id: AbortSignalId,
    state: Arc<Mutex<AbortState>>,
}

impl AbortSignal {
    /// Consumes this signal and waits for either `f` to complete, or the
    /// signal to fire; whichever happens first.
    pub async fn aborting<T, E, F>(self, f: F) -> Result<T, E>
    where
        F: Future<Output = Result<T, E>>,
        E: From<AbortError>,
    {
        futures::pin_mut!(f);
        match future::select(f, self).await {
            Either::Left((result, _)) => result,
            Either::Right((_, _)) => Err(AbortError)?,
        }
    }
}

impl Clone for AbortSignal {
    fn clone(&self) -> Self {
        AbortSignal {
            id: AbortSignalId::Unreserved,
            state: Arc::clone(&self.state),
        }
    }
}

impl Future for AbortSignal {
    type Output = ();

    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        let mut guard = this.state.lock().unwrap();

        let AbortState::Armed { next_id, wakers } = &mut *guard else {
            // Our controller has been aborted. Forget our ID now,
            // so that we can skip re-locking and cleaning up in
            // `drop` later.
            this.id = AbortSignalId::Unreserved;
            return Poll::Ready(());
        };

        let id = match this.id {
            AbortSignalId::Reserved(id) => id,
            AbortSignalId::Unreserved => {
                // A task is waiting on us for the first time.
                // Reserve our ID.
                let this_id = *next_id;
                *next_id = this_id.checked_add(1).unwrap();
                this.id = AbortSignalId::Reserved(this_id);
                this_id
            }
        };

        wakers
            .entry(id)
            .and_modify(|waker| {
                // The `Waker` docs recommend using `clone_from` to
                // avoid cloning a waker that would wake the same task.
                waker.clone_from(context.waker());
            })
            .or_insert_with(|| context.waker().clone());

        Poll::Pending
    }
}

impl Drop for AbortSignal {
    fn drop(&mut self) {
        if let AbortSignalId::Reserved(id) = self.id {
            // We only need to lock and update the shared state if
            // we have at least one task waiting on us.
            if let AbortState::Armed { wakers, .. } = &mut *self.state.lock().unwrap() {
                wakers.remove(&id);
            }
        }
    }
}

impl Debug for AbortSignal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "AbortSignal {{ id: {:?}, .. }}", self.id)
    }
}

/// Uniquely identifies an [`AbortSignal`].
///
/// ### Notes
///
/// [`AbortSignalId`] takes advantage of Rust's [discriminant elision][1] for
/// option-like enums, so it has the same size and alignment as `usize`.
///
/// [1]: https://rust-lang.github.io/unsafe-code-guidelines/layout/enums.html
#[derive(Clone, Copy, Debug)]
enum AbortSignalId {
    /// Indicates that an ID has not been reserved for this [`AbortSignal`],
    /// because no tasks are currently waiting on it.
    Unreserved,

    /// The ID for an [`AbortSignal`] with one or more tasks waiting on it.
    Reserved(NonZeroUsize),
}

/// State shared between an [`AbortController`] and its [`AbortSignal`]s.
///
/// ### Notes
///
/// Controllers and signals share state using an `Arc<Mutex<AbortState>>`.
/// There are more efficient implementations than a single shared lock,
/// but we're prioritizing readability and correctness at-a-glance over
/// raw performance.
enum AbortState {
    /// The controller has not yet been aborted.
    Armed {
        /// The ID to reserve for the next [`AbortSignalId`].
        next_id: NonZeroUsize,

        /// A map of reserved signal IDs to waiting tasks
        /// that should be woken up on abort.
        wakers: BTreeMap<NonZeroUsize, Waker>,
    },

    /// The controller has been aborted.
    Aborted,
}

#[derive(thiserror::Error, Debug)]
#[error("operation aborted")]
pub struct AbortError;