lunar-lib 0.4.4

use std::{
    panic::AssertUnwindSafe,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
        mpsc::{Sender, channel},
    },
};

use crate::error;

mod cli_progress_bar;
pub use cli_progress_bar::*;

/// A thread safe, thread stable progress bar, abstracted for any frontend
#[derive(Clone)]
pub struct ProgressBar {
    handle: Arc<ProgressHandle>,
}

struct ProgressHandle {
    value: AtomicUsize,
    max: AtomicUsize,
    renderer: Arc<dyn ProgressRenderer>,
    sender: Sender<ProgressEvent>,
}

impl Drop for ProgressHandle {
    fn drop(&mut self) {
        let _ = self.sender.send(ProgressEvent::Exit);
        self.renderer.on_finish();
    }
}

impl ProgressBar {
    /// Runs a custom closure for this progress bar.
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    ///
    /// # Notes
    ///
    /// If the input closure panics, it will not kill the thread, it will instead log the error and return
    pub fn custom<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        self.handle
            .sender
            .send(ProgressEvent::Custom {
                f: Box::new(f),
                _handle: self.handle.clone(),
            })
            .expect("Progress handler thread died.");
    }

    /// Waits until the flush is processed, allowing you to wait until all previous events have been processed
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    pub fn flush(&self) {
        let (tx, rx) = channel();
        self.handle
            .sender
            .send(ProgressEvent::Flush {
                done: tx,
                _handle: self.handle.clone(),
            })
            .expect("Progress handler thread died.");
        let _ = rx.recv();
    }

    /// Gets the max value of the internal `ProgressHandle`
    ///
    /// # Notes
    ///
    /// While this will return the correct max value at time of call, In multi-threaded contexts, this max may have already changed by the time it has been used
    pub fn get_max(&self) -> usize {
        self.handle.max.load(Ordering::SeqCst)
    }

    /// Gets the value of the internal `ProgressHandle`
    ///
    /// # Notes
    ///
    /// While this will return the correct value at time of call, In multi-threaded contexts, this value may have already changed by the time it has been used
    pub fn get_value(&self) -> usize {
        self.handle.value.load(Ordering::SeqCst)
    }

    /// Increments the value of a progress bar by 1
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    pub fn increment(&self) {
        self.handle
            .sender
            .send(ProgressEvent::Increment {
                handle: self.handle.clone(),
            })
            .expect("Progress handler thread died.");
    }

    /// Increments the value of the progress bar by 1 and runs a pre-write and post-write closure
    ///
    /// The `pre` closure, if some, will run before the value is incremented
    /// The `post` closure, if some, will run after the value is incremented, and before [`ProgressRenderer::on_update()`] method is run
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    pub fn increment_and_run<P, Q>(&self, pre: Option<P>, post: Option<Q>)
    where
        P: FnOnce() + Send + 'static,
        Q: FnOnce() + Send + 'static,
    {
        self.handle
            .sender
            .send(ProgressEvent::IncrementAndRun {
                handle: self.handle.clone(),
                pre: pre.map(box_dyn),
                post: post.map(box_dyn),
            })
            .expect("Progress handler thread died.");
    }

    /// Creates a new progress bar instance using the renderer
    ///
    /// This progress bar can safely be shared and updated between threads
    pub fn new(value: usize, max: usize, renderer: Arc<dyn ProgressRenderer>) -> Self {
        let (tx, rx) = channel();

        let handle = Arc::new(ProgressHandle {
            value: AtomicUsize::new(value),
            max: AtomicUsize::new(max),
            renderer: renderer.clone(),
            sender: tx,
        });

        spawn_progress_thread(rx);

        renderer.on_start(value, max);

        Self { handle }
    }

    /// Sets the max-value of the progress bar
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    pub fn set_max(&self, max: usize) {
        self.handle
            .sender
            .send(ProgressEvent::SetMax {
                handle: self.handle.clone(),
                max,
            })
            .expect("Progress handler thread died.");
    }

    /// Sets the max-value of the progress bar and runs a pre-write and post-write closure
    ///
    /// The `pre` closure, if some, will run before the value is incremented
    /// The `post` closure, if some, will run after the value is incremented, and before [`ProgressRenderer::on_update()`] method is run
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    ///
    /// # Notes
    ///
    /// If the any of the input closures panic, it will not kill the thread, it will instead log the error and ignore it
    pub fn set_max_and_run<P, Q>(&self, max: usize, pre: Option<P>, post: Option<Q>)
    where
        P: FnOnce() + Send + 'static,
        Q: FnOnce() + Send + 'static,
    {
        self.handle
            .sender
            .send(ProgressEvent::SetMaxAndRun {
                handle: self.handle.clone(),
                max,
                pre: pre.map(box_dyn),
                post: post.map(box_dyn),
            })
            .expect("Progress handler thread died.");
    }

    /// Sets the value of the progress bar
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    pub fn set_value(&self, value: usize) {
        self.handle
            .sender
            .send(ProgressEvent::SetValue {
                handle: self.handle.clone(),
                value,
            })
            .expect("Progress handler thread died.");
    }

    /// Sets the value of the progress bar and runs a pre-write and post-write closure
    ///
    /// The `pre` closure, if some, will run before the value is incremented
    /// The `post` closure, if some, will run after the value is incremented, and before [`ProgressRenderer::on_update()`] method is run
    ///
    /// # Panics
    ///
    /// Panics if the internal `GLOBAL_SENDER` thread cannot be communicated with. This usually signals critical failure somewhere
    ///
    /// # Notes
    ///
    /// If the any of the input closures panic, it will not kill the thread, it will instead log the error and ignore it
    pub fn set_value_and_run<P, Q>(&self, value: usize, pre: Option<P>, post: Option<Q>)
    where
        P: FnOnce() + Send + 'static,
        Q: FnOnce() + Send + 'static,
    {
        self.handle
            .sender
            .send(ProgressEvent::SetValueAndRun {
                handle: self.handle.clone(),
                value,
                pre: pre.map(box_dyn),
                post: post.map(box_dyn),
            })
            .expect("Progress handler thread died.");
    }

    /// Forces the [`ProgressRenderer`] to upate with the new values
    ///
    /// # Panics
    ///
    /// Panics if the progress bars handling thread cannot be communicated with.
    pub fn update(&self) {
        self.handle
            .sender
            .send(ProgressEvent::Update {
                handle: self.handle.clone(),
            })
            .expect("Progress handler thread died.")
    }

    /// Checks if the the handling thread is alive by sending a check to it and seeing if it returns [`Ok`]
    ///
    /// # Notes
    ///
    /// This function will correctly check if the handling thread is alive at the time of calling, but can not guarantee the thread will stay alive after calling or that it receives the check
    pub fn is_alive(&self) -> bool {
        self.handle.sender.send(ProgressEvent::Check).is_ok()
    }

    /// Calls the [`set_label()`] function for the held [`ProgressRenderer`]
    pub fn set_label(&self, label: &str) {
        self.handle.renderer.set_label(label);
    }

    /// Calls the [`on_notify()`] function for the held [`ProgressRenderer`] and then instantly updates self
    pub fn notify(&self, message: impl AsRef<str>) {
        self.handle.renderer.on_notify(message.as_ref());
        self.update();
    }
}

/// Defines how a progress bar should be rendered
pub trait ProgressRenderer: Send + Sync {
    /// Called when the [`ProgressBar`] which holds the instance of [`Self`] is spawned
    fn on_start(&self, value: usize, max: usize);

    /// Called when the [`ProgressBar`] which holds the instance of [`Self`] is updated
    fn on_update(&self, value: usize, max: usize);

    /// Called when the [`ProgressBar`] which holds the instance of [`Self`] is dropped
    ///
    /// # Notes
    ///
    /// [`Self::on_update()`] is NOT automatically called before finish, if you would like to update before the renderer is dropped, you can do so when defining this function
    fn on_finish(&self);

    /// Generic notification. Can be sent from anything with access to the [`ProgressBar`] that holds the instance of [`Self`]
    ///
    /// This function is not called internally by anything, its meant for other libraries to call when they want to notify your renderer of some message. If you do not care, you can ignore this function
    fn on_notify(&self, msg: &str);

    /// Sets the label of [`Self`]. Can be sent from anything with access to the [`ProgressBar`] that holds the instance of [`Self`]
    ///
    /// This function is not called internally by anything, its meant for other libraries to call when they want to label your progress bar. If you do not care, you can ignore this function
    fn set_label(&self, msg: &str);
}

// ==============================================================
// Thread stuff
// ==============================================================

enum ProgressEvent {
    Check,
    Exit,
    Update {
        handle: Arc<ProgressHandle>,
    },
    Increment {
        handle: Arc<ProgressHandle>,
    },
    IncrementAndRun {
        handle: Arc<ProgressHandle>,
        pre: Option<Box<dyn FnOnce() + Send>>,
        post: Option<Box<dyn FnOnce() + Send>>,
    },
    SetValue {
        handle: Arc<ProgressHandle>,
        value: usize,
    },
    SetValueAndRun {
        handle: Arc<ProgressHandle>,
        value: usize,
        pre: Option<Box<dyn FnOnce() + Send>>,
        post: Option<Box<dyn FnOnce() + Send>>,
    },
    SetMax {
        handle: Arc<ProgressHandle>,
        max: usize,
    },
    SetMaxAndRun {
        handle: Arc<ProgressHandle>,
        max: usize,
        pre: Option<Box<dyn FnOnce() + Send>>,
        post: Option<Box<dyn FnOnce() + Send>>,
    },
    Custom {
        _handle: Arc<ProgressHandle>,
        f: Box<dyn FnOnce() + Send>,
    },
    Flush {
        _handle: Arc<ProgressHandle>,
        done: std::sync::mpsc::Sender<()>,
    },
}

fn spawn_progress_thread(rx: std::sync::mpsc::Receiver<ProgressEvent>) {
    std::thread::spawn(move || {
        while let Ok(event) = rx.recv() {
            match event {
                ProgressEvent::Check => {}
                ProgressEvent::Exit => {
                    break;
                }
                ProgressEvent::Increment { handle } => {
                    let value = handle.value.fetch_add(1, Ordering::SeqCst);
                    let max = handle.max.load(Ordering::SeqCst);
                    handle.renderer.on_update(value + 1, max);
                }
                ProgressEvent::IncrementAndRun { handle, pre, post } => {
                    run_optional_log_panic(pre);
                    let value = handle.value.fetch_add(1, Ordering::SeqCst);
                    run_optional_log_panic(post);
                    let max = handle.max.load(Ordering::SeqCst);
                    handle.renderer.on_update(value + 1, max);
                }
                ProgressEvent::SetValue { handle, value } => {
                    handle.value.store(value, Ordering::SeqCst);
                    let max = handle.max.load(Ordering::SeqCst);
                    handle.renderer.on_update(value, max);
                }
                ProgressEvent::SetValueAndRun {
                    handle,
                    value,
                    pre,
                    post,
                } => {
                    handle.value.store(value, Ordering::SeqCst);
                    run_optional_log_panic(pre);
                    let max = handle.max.load(Ordering::SeqCst);
                    run_optional_log_panic(post);
                    handle.renderer.on_update(value, max);
                }
                ProgressEvent::SetMax { max, handle } => {
                    let value = handle.value.load(Ordering::SeqCst);
                    handle.max.store(max, Ordering::SeqCst);
                    handle.renderer.on_update(value, max);
                }
                ProgressEvent::SetMaxAndRun {
                    handle,
                    max,
                    pre,
                    post,
                } => {
                    let value = handle.value.load(Ordering::SeqCst);
                    run_optional_log_panic(pre);
                    handle.max.store(max, Ordering::SeqCst);
                    run_optional_log_panic(post);
                    handle.renderer.on_update(value, max);
                }
                ProgressEvent::Custom { f, _handle } => run_or_log_panic(f),
                ProgressEvent::Flush { done, _handle } => done.send(()).unwrap(),
                ProgressEvent::Update { handle } => {
                    let value = handle.value.load(Ordering::SeqCst);
                    let max = handle.max.load(Ordering::SeqCst);
                    handle.renderer.on_update(value, max);
                }
            }
        }
    });
}

fn box_dyn<F: FnOnce() + Send + 'static>(f: F) -> Box<dyn FnOnce() + Send + 'static> {
    Box::new(f)
}

#[inline(always)]
fn run_or_log_panic<F: FnOnce() + Send + 'static>(f: F) {
    if let Err(err) = std::panic::catch_unwind(AssertUnwindSafe(f)) {
        error!("ProgressBar closure panicked: {err:?}")
    }
}

#[inline(always)]
fn run_optional_log_panic<F: FnOnce() + Send + 'static>(f: Option<F>) {
    if let Some(f) = f {
        run_or_log_panic(f);
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    };

    use crate::progress::{ProgressBar, ProgressRenderer};

    struct TestRenderer {
        log: Arc<Mutex<Vec<String>>>,
    }

    impl TestRenderer {
        fn new() -> Self {
            Self {
                log: Arc::new(Mutex::new(Vec::new())),
            }
        }

        fn log(&self) -> Vec<String> {
            self.log.lock().unwrap().clone()
        }

        fn push_log(&self, log: String) {
            self.log.lock().unwrap().push(log);
        }
    }

    impl ProgressRenderer for TestRenderer {
        fn on_start(&self, value: usize, max: usize) {
            self.push_log(format!("Start: [{value} / {max}]"));
        }

        fn on_update(&self, value: usize, max: usize) {
            self.push_log(format!("Update: [{value} / {max}]"));
        }

        fn on_finish(&self) {
            self.push_log(format!("Finish"));
        }

        fn on_notify(&self, msg: &str) {
            self.push_log(msg.to_owned());
        }

        fn set_label(&self, _msg: &str) {}
    }

    #[test]
    fn progress_increment() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 3, renderer.clone());

        bar.increment();
        bar.increment();
        bar.flush();

        let log = renderer.log();
        assert_eq!(log.len(), 3);
        assert_eq!(log[0], "Start: [0 / 3]");
        assert_eq!(log[1], "Update: [1 / 3]");
        assert_eq!(log[2], "Update: [2 / 3]");
    }

    #[test]
    fn progress_notify() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 3, renderer.clone());

        bar.increment();
        bar.flush();
        bar.notify("notification");
        bar.flush();

        let log = renderer.log();

        log.iter().for_each(|s| println!("{s}"));

        assert_eq!(log.len(), 4);
        assert_eq!(log[0], "Start: [0 / 3]");
        assert_eq!(log[1], "Update: [1 / 3]");
        assert_eq!(log[2], "notification");
        assert_eq!(log[3], "Update: [1 / 3]");
    }

    #[test]
    fn progress_finishes_on_drop() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 0, renderer.clone());

        drop(bar);

        let log = renderer.log();
        assert_eq!(log.len(), 2);
        assert_eq!(log[0], "Start: [0 / 0]");
        assert_eq!(log[1], "Finish")
    }

    #[test]
    fn progress_closures_run_on_increment() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 1, renderer.clone());

        let pre_ran = Arc::new(AtomicBool::new(false));
        let post_ran = Arc::new(AtomicBool::new(false));

        let pre_flag = pre_ran.clone();
        let post_flag = post_ran.clone();

        bar.increment_and_run(
            Some(move || pre_flag.store(true, Ordering::SeqCst)),
            Some(move || post_flag.store(true, Ordering::SeqCst)),
        );
        bar.flush();

        assert!(pre_ran.load(Ordering::SeqCst));
        assert!(post_ran.load(Ordering::SeqCst));

        let log = renderer.log();
        assert_eq!(log.len(), 2);
        assert_eq!(log[0], "Start: [0 / 1]");
        assert_eq!(log[1], "Update: [1 / 1]");
    }

    #[test]
    fn progress_closures_run_on_set_value() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 5, renderer.clone());

        let pre_ran = Arc::new(AtomicBool::new(false));
        let post_ran = Arc::new(AtomicBool::new(false));

        let pre_flag = pre_ran.clone();
        let post_flag = post_ran.clone();

        bar.set_value_and_run(
            5,
            Some(move || pre_flag.store(true, Ordering::SeqCst)),
            Some(move || post_flag.store(true, Ordering::SeqCst)),
        );
        bar.flush();

        assert!(pre_ran.load(Ordering::SeqCst));
        assert!(post_ran.load(Ordering::SeqCst));

        let log = renderer.log();
        assert_eq!(log.len(), 2);
        assert_eq!(log[0], "Start: [0 / 5]");
        assert_eq!(log[1], "Update: [5 / 5]");
    }

    #[test]
    fn progress_closures_run_on_set_max() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 3, renderer.clone());

        let pre_ran = Arc::new(AtomicBool::new(false));
        let post_ran = Arc::new(AtomicBool::new(false));

        let pre_flag = pre_ran.clone();
        let post_flag = post_ran.clone();

        bar.set_max_and_run(
            5,
            Some(move || pre_flag.store(true, Ordering::SeqCst)),
            Some(move || post_flag.store(true, Ordering::SeqCst)),
        );
        bar.flush();

        assert!(pre_ran.load(Ordering::SeqCst));
        assert!(post_ran.load(Ordering::SeqCst));

        let log = renderer.log();
        assert_eq!(log.len(), 2);
        assert_eq!(log[0], "Start: [0 / 3]");
        assert_eq!(log[1], "Update: [0 / 5]");
    }

    #[test]
    fn progress_thread_safe_on_closure_panic() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 100, renderer.clone());

        bar.increment_and_run(Some(|| panic!()), Some(|| ()));
        bar.flush();
    }

    #[test]
    fn progress_thread_stability() {
        let renderer = Arc::new(TestRenderer::new());
        let bar = ProgressBar::new(0, 100, renderer.clone());

        let handles: Vec<_> = (0..5)
            .map(|_| {
                let bar = bar.clone();
                std::thread::spawn(move || {
                    for _ in 0..20 {
                        bar.increment();
                    }
                })
            })
            .collect();

        for h in handles {
            h.join().unwrap()
        }

        bar.flush();

        let log = renderer.log();
        assert_eq!(log.len(), 101);

        assert_eq!(log[0], "Start: [0 / 100]");
        for (i, entry) in log.iter().skip(1).enumerate() {
            assert_eq!(entry, &format!("Update: [{i} / 100]", i = i + 1))
        }
    }
}