ufotofu 0.12.5

use crate::prelude::*;

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

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

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

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

/// A wrapper for a [`Producer`] which allows concurrent read access to the [`Status`] of the wrapped producer, and the number of [`items`](Producer::Item) produced 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 producer.
///
/// Created via the [`ProducerExt::to_watch`] method.
///
/// <br/>Counterpart: the [`consumer::Watch`] type.
pub struct Watch<P> {
    inner: P,
    watcher: Watcher,
}

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

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

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

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

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

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

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

impl<P> Producer for Watch<P>
where
    P: Producer,
{
    type Item = P::Item;

    type Final = P::Final;

    type Error = P::Error;

    async fn produce(&mut self) -> Result<Either<Self::Item, Self::Final>, Self::Error> {
        match self.inner.produce().await {
            Ok(Left(item)) => {
                self.watcher.0.count.update(|n| {
                    n.checked_add(1)
                        .expect("producer watcher cannot count more than usize::MAX items")
                });
                Ok(Left(item))
            }
            Ok(Right(fin)) => {
                self.watcher.0.status.set(Status::Finalised);
                Ok(Right(fin))
            }
            Err(err) => {
                self.watcher.0.status.set(Status::Errored);
                Err(err)
            }
        }
    }

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

impl<P> BulkProducer for Watch<P>
where
    P: BulkProducer,
{
    async fn expose_items_gracefully<F, R>(
        &mut self,
        f: F,
    ) -> Result<Either<R, (F, Self::Final)>, (F, Self::Error)>
    where
        F: AsyncFnOnce(&[Self::Item]) -> (usize, R),
    {
        let mut f = Some(f);

        match self
            .inner
            .expose_items_gracefully(async |items| {
                let f = f.take().expect("constructed as a Some variant");
                let (produced, result) = f(items).await;
                self.watcher.0.count.update(|n| {
                    n.checked_add(produced)
                        .expect("producer watcher cannot count more than usize::MAX items")
                });
                (produced, result)
            })
            .await
        {
            Ok(Left(result)) => Ok(Left(result)),
            Ok(Right((_, fin))) => {
                self.watcher.0.status.set(Status::Finalised);
                Ok(Right((
                    f.take().expect(
                        "provided closure must not be called when the final item is encountered",
                    ),
                    fin,
                )))
            }
            Err((_, err)) => {
                self.watcher.0.status.set(Status::Errored);
                Err((
                    f.take()
                        .expect("provided closure must not be called when an error occurs"),
                    err,
                ))
            }
        }
    }
}