ufotofu 0.12.5

use crate::prelude::*;

use alloc::rc::Rc;
use core::cell::Cell;

/// The possible states for a [`Consumer`] to be in, as reported by, e.g., a [`Watcher`] or a [`Stats`](consumer::Stats) wrapper.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Status {
    /// The state of a [`Consumer`] which may continue to consume [`items`](Consumer::Item).
    Processing,
    /// The state of a [`Consumer`] which has emitted a [`Final`](Consumer::Final) value, and will therefore consume no further [`items`](Consumer::Item).
    Finalised,
    /// The state of a [`Consumer`] which has reported an [`Error`](Consumer::Error), and will therefore consume no further [`items`](Consumer::Item).
    Errored,
}

struct WatcherInner {
    count: Cell<usize>,
    status: Cell<Status>,
}

/// Provides concurrent access to basic information about a [`watched`](Watch) [`Consumer`].
///
/// In particular, it is possible to query a count of the number of [`items`](Consumer::Item) consumed, and the [`Status`] of the consumer.
///
/// A `Watcher` can be cheaply cloned to provide access to this information in multiple places.
///
/// Created via the [`Watch::watcher`] method.
///
/// <br/>Counterpart: the [`producer::Watcher`] type.
#[derive(Clone)]
pub struct Watcher(Rc<WatcherInner>);

/// A wrapper for a [`Consumer`] which allows concurrent read access to the [`Status`] of the wrapped consumer, and the number of [`items`](Consumer::Item) consumed since the wrapper was created.
///
/// The [`watcher`](Watch::watcher) method returns a [`Watcher`] through which to read the count and status. The [`count`](Watch::count) method can be used to access the count on the wrapper directly.
///
/// Use the [`into_inner`](Watch::into_inner) method to consume the wrapper and access the wrapped consumer.
///
/// Created via the [`ConsumerExt::to_watch`] method.
///
/// <br/>Counterpart: the [`producer::Watch`] type.
pub struct Watch<C> {
    inner: C,
    watcher: Watcher,
}

impl Watcher {
    fn new() -> Self {
        Self(Rc::new(WatcherInner {
            count: Cell::new(0),
            status: Cell::new(Status::Processing),
        }))
    }

    /// Queries the number of [`items`](Consumer::Item) produced by the [`Consumer`] [`watched`](Watch) by this `Watcher` since it started being watched.
    pub fn count(&self) -> usize {
        self.0.count.get()
    }

    /// Queries the [`Status`] of the [`Consumer`] [`watched`](Watch) by this `Watcher`.
    pub fn status(&self) -> Status {
        self.0.status.get()
    }
}

impl<C> Watch<C> {
    pub(crate) fn new(inner: C) -> Self {
        let watcher = Watcher::new();
        Watch { inner, watcher }
    }

    /// Creates a new [`Watcher`] for information about the wrapped [`Consumer`].
    pub fn watcher(&self) -> Watcher {
        self.watcher.clone()
    }

    /// Returns the number of [`items`](Consumer::Item) which have been consumed since the wrapped consumer started being watched.
    pub fn count(&self) -> usize {
        self.watcher.count()
    }

    /// Consumes the wrapper to return the inner [`Consumer`].
    pub fn into_inner(self) -> C {
        self.inner
    }
}

impl<C> Consumer for Watch<C>
where
    C: Consumer,
{
    type Item = C::Item;

    type Final = C::Final;

    type Error = C::Error;

    async fn consume(&mut self, val: Either<Self::Item, Self::Final>) -> Result<(), Self::Error> {
        let finalised = val.is_right();
        match self.inner.consume(val).await {
            Ok(()) => {
                if finalised {
                    self.watcher.0.status.set(Status::Finalised);
                } else {
                    self.watcher.0.count.update(|n| {
                        n.checked_add(1)
                            .expect("consumer watcher cannot count more than usize::MAX items")
                    });
                }
                Ok(())
            }
            Err(err) => {
                self.watcher.0.status.set(Status::Errored);
                Err(err)
            }
        }
    }

    async fn flush(&mut self) -> Result<(), Self::Error> {
        let result = self.inner.flush().await;
        if result.is_err() {
            self.watcher.0.status.set(Status::Errored)
        }
        result
    }
}

impl<C> BulkConsumer for Watch<C>
where
    C: BulkConsumer,
{
    async fn expose_slots_gracefully<F, R>(&mut self, f: F) -> Result<R, (F, Self::Error)>
    where
        F: AsyncFnOnce(&mut [Self::Item]) -> (usize, R),
    {
        let mut f = Some(f);
        match self
            .inner
            .expose_slots_gracefully(async |items| {
                let f = f.take().expect("constructed as a Some variant");
                let (consumed, result) = f(items).await;
                self.watcher.0.count.update(|n| {
                    n.checked_add(consumed)
                        .expect("consumer watcher cannot count more than usize::MAX items")
                });
                (consumed, result)
            })
            .await
        {
            Ok(result) => Ok(result),
            Err((_, err)) => {
                self.watcher.0.status.set(Status::Errored);
                Err((
                    f.take()
                        .expect("provided closure must not be called when an error occurs"),
                    err,
                ))
            }
        }
    }
}