lolistealer/src/work_async/progress.rs

use super::*;
use tokio::{
    sync::{
	mpsc::{
	    channel,
	    unbounded_channel,
	    Sender,
	    Receiver,
	    UnboundedReceiver,
	    UnboundedSender
	},
	Mutex,
    },
    task::{
	self,
	JoinHandle,
    },
};

use termprogress::{
    spinner,
    Display,
    Spinner,
};

use std::{
    sync::Arc,
    fmt,
};

enum CommandInternal
{
    PrintLine(String),
    Bump,
    Kill(Option<String>),
    PushTask(String),
    PopTask(String),
}

// TODO: Change from `Spinner` to `Bar`. Have push max, push current, etc.

struct Command
{
    internal: CommandInternal,
    callback: UnboundedSender<()>,
}

impl std::ops::Drop for Command
{
    fn drop(&mut self)
    {
	let _ = self.callback.send(());
    }
}

pub struct AsyncProgressCounter
{
    writer: Sender<Command>,
    reader: Receiver<Command>,
    title: String,
}

#[derive(Clone)]
pub struct CommandSender(Sender<Command>, Option<String>);

pub struct CommandCallback(UnboundedReceiver<()>);

pub mod error;

impl Command
{
    pub fn new(internal: CommandInternal) -> (Self, CommandCallback)
    {
	let (tx, rx) = unbounded_channel();
	(Self {
	    internal,
	    callback: tx,
	}, CommandCallback(rx))
    }
}

macro_rules! prog_send {
    (try link unwind $expression:expr) => {
	{
	    $expression.await.expect("mpsc fatal").wait().await.expect("mpsc fatal")
	}
    };
    (try link $expression:expr) => {
	{
	    $expression.await?.wait().await?
	}
    };
    (try unwind $expression:expr) => {
	{
	    let _ =  $expression.await.expect("mpsc fatal");
	}
    };
    (try $expression:expr) => {
	{
	    let _ =$expression.await?;
	}
    };
    (link $expression:expr) => {
	{
	    if let Ok(ok) = $expression.await {
		let _ = ok.wait().await;
	    }
	}
    };
    ($expression:expr) => {
	{
	    let _ = $expression.await;
	}
    };
}

impl CommandSender
{
    /// Clone with a new opt string
    pub fn clone_with(&self, opt: impl Into<String>) -> Self
    {
	Self (
	    self.0.clone(),
	    Some(opt.into())
	)
    }
    
    fn opt_str(&self, string: impl AsRef<str>) -> String
    {
	if let Some(opt) = &self.1 {
	    format!("[{}] {}", opt, string.as_ref())
	} else {
	    string.as_ref().into()
	}
    }

    /// Get or set an opt string
    pub fn opt(&mut self) -> &mut Option<String>
    {
	&mut self.1
    }
    
    /// Print a line
    pub async fn println(&mut self, string: impl AsRef<str>) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::PrintLine(self.opt_str(string)));
	self.0.send(com).await?;
	Ok(call)
    }

    /// Finalise the progress counter
    pub async fn bump(&mut self) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::Bump);
	self.0.send(com).await?;
	Ok(call)
    }
    
    /// Finalise the progress counter
    pub async fn kill(&mut self) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::Kill(None));
	self.0.send(com).await?;
	Ok(call)
    }

    /// Finalise the progress counter and print a line
    pub async fn kill_with(&mut self, string: impl AsRef<str>) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::Kill(Some(self.opt_str(string))));
	self.0.send(com).await?;
	Ok(call)
    }


    /// Remove a task
    pub async fn pop_task(&mut self, string: impl Into<String>) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::PopTask(string.into()));
	self.0.send(com).await?;
	Ok(call)
    }

    /// Add a task
    pub async fn push_task(&mut self, string: impl Into<String>) -> Result<CommandCallback, error::Error>
    {
	let (com, call) = Command::new(CommandInternal::PushTask(string.into()));
	self.0.send(com).await?;
	Ok(call)
    }
}

impl CommandCallback
{
    /// Wait until the command has been processed.
    pub async fn wait(mut self) -> Result<(), error::Error>
    {
	self.0.recv().await.ok_or(error::Error::RecvError)
    }
}

impl AsyncProgressCounter
{
    /// Create a new `AsyncProgressCounter`
    pub fn new(title: impl Into<String>) -> Self
    {
	let (tx, rx) = channel(16);

	Self {
	    reader: rx,
	    writer: tx,
	    title: title.into(),
	}
    }
    
    /// Create a new writer with an opt string
    pub fn writer_with(&self, opt: impl Into<String>) -> CommandSender
    {
	CommandSender(self.writer.clone(), Some(opt.into()))
    }
    /// Create a new writer
    pub fn writer(&self) -> CommandSender
    {
	CommandSender(self.writer.clone(), None)
    }

    /// Consume the instance and host it's receiver
    pub fn host(mut self) -> JoinHandle<()>
    {
	let mut spin = spinner::Spin::with_title(&self.title[..], Default::default());

	let mut task = tasklist::TaskList::new();
	
	task::spawn(async move {
	    while let Some(com) = self.reader.recv().await {
		match &com.internal {
		    CommandInternal::PrintLine(line) => spin.println(&line[..]),
		    CommandInternal::Bump => {
			spin.bump(); //TODO: We'll have to change this when we change to real progress
		    },
		    CommandInternal::Kill(Some(line)) => {
			spin.complete_with(line.as_str());
			self.reader.close();
			break;
		    },
		    CommandInternal::Kill(_) => {
			spin.complete();
			self.reader.close();
			break;
		    },
		    CommandInternal::PushTask(tstr) => {
			task.push(tstr);
			spin.set_title(task.as_str());
		    },
		    CommandInternal::PopTask(tstr) => {
			task.pop_value(tstr);
			spin.set_title(task.as_str());
		    },
		}
	    }
	})
    }
}