osc94 0.1.3

//! This module provides an easy adapter for iterators.

use super::{OSC94, ProgressState};
use std::{
    io::{Error as IoError, Stderr, Write},
    iter::FusedIterator,
};

/// An iterator adapter that writes OSC 9;4 progress updates while iterating.
///
/// ## Construction
///
/// Can be created with iterator adapter [`ProgressIteratorExt::with_progress`], or using the [`ProgressIter::new`] constructor for explicity.
///
/// ## Methods
///
/// - [`ProgressIter::total`]: override the total item count used to calculate progress percentage.
///     - By default, the adapter will attempt to determine an exact total from the underlying iterator's [`size_hint`](Iterator::size_hint)
///     - If no exact total is known, the adapter uses the [indeterminate state](ProgressState::Indeterminate).
/// - [`ProgressIter::get_total`]: returns the total item count used to calculate progress percentage.
/// - [`ProgressIter::get_current`]: returns the number of items yielded so far.
/// - [`ProgressIter::last_error`]: returns the last IO error produced while writing progress updates, if any.
#[derive(Debug)]
pub struct ProgressIter<I, W: Write = Stderr> {
    iter: I,
    destination: W,
    total: Option<usize>,
    current: usize,
    started: bool,
    last_error: Option<IoError>,
}

impl<I, W: Write> ProgressIter<I, W>
where
    I: Iterator,
{
    /// Creates a new progress iterator that writes to stderr.
    pub fn new(iter: I, destination: W) -> Self {
        let total = exact_total(&iter);

        Self {
            iter,
            destination,
            total,
            current: 0,
            started: false,
            last_error: None,
        }
    }
}

impl<I, W> ProgressIter<I, W>
where
    I: Iterator,
    W: Write,
{
    /// Overrides the total item count used to calculate progress percentage.
    #[must_use]
    pub const fn total(mut self, total: usize) -> Self {
        self.total = Some(total);
        self
    }

    /// Returns the total item count used to calculate progress percentage.
    pub const fn get_total(&self) -> Option<usize> {
        self.total
    }

    /// Returns the number of items yielded so far.
    pub const fn get_current(&self) -> usize {
        self.current
    }

    /// Returns the last IO error produced while writing progress updates.
    pub const fn last_error(&self) -> Option<&IoError> {
        self.last_error.as_ref()
    }

    fn start(&mut self) {
        if self.started {
            return;
        }

        self.started = true;

        match self.total {
            Some(_) => self.flush(ProgressState::Normal, 0),
            None => self.flush(ProgressState::Indeterminate, 0),
        }
    }

    fn update(&mut self) {
        let Some(total) = self.total else {
            return;
        };

        let progress = percentage(self.current, total);
        self.flush(ProgressState::Normal, progress);
    }

    fn flush(&mut self, state: ProgressState, progress: u8) {
        let raw = OSC94 { state, progress };

        if let Err(error) = write!(self.destination, "{raw}") {
            self.last_error = Some(error);
        }
    }
}

impl<I, W> Iterator for ProgressIter<I, W>
where
    I: Iterator,
    W: Write,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.iter.next()?;

        self.start();
        self.current = self.current.saturating_add(1);
        self.update();

        Some(item)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<I, W: Write> Drop for ProgressIter<I, W> {
    fn drop(&mut self) {
        let raw = OSC94::default();
        let _ = write!(self.destination, "{raw}");
    }
}

impl<I, W> FusedIterator for ProgressIter<I, W>
where
    I: FusedIterator,
    W: Write,
{
}

/// Extension trait for wrapping iterators with OSC 9;4 progress updates.
pub trait ProgressIteratorExt: Iterator + Sized {
    /// Wraps this iterator with automatic progress updates. See [`ProgressIter`] for details.
    fn with_progress(self) -> ProgressIter<Self> {
        self.with_progress_to(std::io::stderr())
    }
    /// Wraps this iterator with automatic progress updates, writing to the specified destination.
    fn with_progress_to<W: Write>(self, destination: W) -> ProgressIter<Self, W>;
}

impl<I> ProgressIteratorExt for I
where
    I: Iterator,
{
    fn with_progress_to<W>(self, destination: W) -> ProgressIter<Self, W>
    where
        W: Write,
    {
        ProgressIter::new(self, destination)
    }
}

fn exact_total<I: Iterator>(iter: &I) -> Option<usize> {
    let (lower, upper) = iter.size_hint();
    upper.filter(|upper| *upper == lower)
}

fn percentage(current: usize, total: usize) -> u8 {
    if total == 0 {
        return 100;
    }

    let progress = current.saturating_mul(100) / total;
    u8::try_from(progress.min(100)).expect("progress is capped at 100")
}