use std::future::Future;
use std::task::{Context, Poll};
use std::time::Instant;
use std::{io::Write, pin::Pin};
use crossterm::{QueueableCommand, cursor, terminal};
use futures::stream::FuturesUnordered;
use futures_lite::{FutureExt as _, Stream};
use owo_colors::OwoColorize;
use crate::spinner::{Spinner, Ticks};
use crate::term::clear_line;
struct Annotated<F> {
inner: F,
id: usize,
}
impl<F> Annotated<F>
where
F: Future,
{
fn new(inner: F, id: usize) -> Self {
Self { inner, id }
}
}
impl<F> Future for Annotated<F>
where
F: Future + Unpin,
{
type Output = (F::Output, usize);
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.get_mut();
let id = this.id;
match this.inner.poll(cx) {
Poll::Ready(output) => Poll::Ready((output, id)),
Poll::Pending => Poll::Pending,
}
}
}
pub struct Monitored<'a, F> {
inner: FuturesUnordered<Annotated<F>>,
ticks: Ticks<'a>,
annotations: Vec<Option<String>>,
annotation_style: owo_colors::Style,
spinner: Option<char>,
spinner_style: owo_colors::Style,
with_elapsed_time: bool,
start: Option<Instant>,
}
impl<'a, F> Monitored<'a, F>
where
F: Future,
{
pub fn new(spinner: Spinner<'a>) -> Self {
Self {
inner: FuturesUnordered::new(),
ticks: spinner.ticks(),
annotations: Vec::new(),
annotation_style: owo_colors::Style::new(),
spinner: None,
spinner_style: owo_colors::Style::new(),
with_elapsed_time: false,
start: None,
}
}
pub fn with_spinner_style(mut self, spinner_style: owo_colors::Style) -> Self {
self.spinner_style = spinner_style;
self
}
pub fn with_annotation_style(mut self, annotation_style: owo_colors::Style) -> Self {
self.annotation_style = annotation_style;
self
}
pub fn with_elapsed_time(mut self, with_elapsed_time: bool) -> Self {
self.with_elapsed_time = with_elapsed_time;
self
}
pub fn push(&mut self, fut: F, annotation: String) {
let id = self.annotations.len();
self.annotations.push(Some(annotation));
self.inner.push(Annotated::new(fut, id));
}
}
impl<F> Stream for Monitored<'_, F>
where
F: Future + Unpin,
{
type Item = F::Output;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
let inner = Pin::new(&mut this.inner);
let ticks = Pin::new(&mut this.ticks);
let elapsed = this.start.get_or_insert_with(|| Instant::now()).elapsed();
if let Poll::Ready(spinner) = ticks.poll_next(cx) {
this.spinner = spinner;
}
let mut stdout = std::io::stdout();
let _ = stdout.queue(cursor::Hide);
for annotation in &this.annotations {
let _ = clear_line(&mut stdout);
if let Some(spinner) = &this.spinner {
print!("{} ", spinner.style(this.spinner_style));
}
if this.with_elapsed_time {
print!("[{:.2}s] ", elapsed.as_secs_f64());
}
if let Some(annotation) = annotation {
println!("{}", annotation.style(this.annotation_style));
}
}
let item = match inner.poll_next(cx) {
Poll::Ready(Some((output, id))) => {
this.annotations[id] = None;
let _ = remove_last_line(&mut stdout);
Poll::Ready(Some(output))
}
Poll::Ready(None) => {
let _ = reset(&mut stdout);
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
};
if !matches!(item, Poll::Ready(None)) {
let active_futures = this.annotations.iter().filter(|a| a.is_some()).count();
if active_futures > 0 {
let _ = stdout.queue(cursor::MoveUp(active_futures as u16));
}
}
let _ = stdout.flush();
item
}
}
fn remove_last_line(stdout: &mut std::io::Stdout) -> std::io::Result<()> {
stdout
.queue(terminal::Clear(terminal::ClearType::CurrentLine))?
.queue(cursor::MoveUp(1))?
.queue(terminal::Clear(terminal::ClearType::CurrentLine))?;
Ok(())
}
fn reset(stdout: &mut std::io::Stdout) -> std::io::Result<()> {
stdout
.queue(cursor::Show)?
.queue(terminal::Clear(terminal::ClearType::CurrentLine))?
.queue(cursor::MoveToColumn(0))?;
Ok(())
}