deja-vu/core/src/tasks.rs

// This file is part of Deja-vu.
//
// Deja-vu is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Deja-vu is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
// See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Deja-vu. If not, see <https://www.gnu.org/licenses/>.
//

use std::{
    collections::VecDeque,
    path::PathBuf,
    process::{Command, Stdio},
    sync::{Arc, Mutex},
    thread::{self, sleep, JoinHandle},
    time::Duration,
};

use color_eyre::Result;
use tokio::sync::mpsc;
use tracing::info;

use crate::{
    action::Action,
    components::popup,
    system::{disk, diskpart},
};

#[derive(Clone, Debug)]
pub enum Task {
    Command(PathBuf, Vec<String>), // (command, args)
    Diskpart(String),              // (script_as_string)
    ScanDisks,
    Sleep,
    UpdateDestDisk(usize), // (disk_index)
    UpdateDiskList,
}

#[derive(Debug)]
pub struct Tasks {
    action_tx: mpsc::UnboundedSender<Action>,
    disk_list: Arc<Mutex<Vec<disk::Disk>>>,
    handle: Option<JoinHandle<()>>,
    task_list: VecDeque<Task>,
    task_rx: mpsc::UnboundedReceiver<Action>, // Used to forward Actions from Tasks to App
    task_tx: mpsc::UnboundedSender<Action>,   // Used to forward Actions from Tasks to App
}

impl Tasks {
    pub fn new(
        action_tx: mpsc::UnboundedSender<Action>,
        disk_list_arc: Arc<Mutex<Vec<disk::Disk>>>,
    ) -> Self {
        let (task_tx, task_rx) = mpsc::unbounded_channel();
        Tasks {
            action_tx,
            disk_list: disk_list_arc,
            handle: None,
            task_list: VecDeque::new(),
            task_rx,
            task_tx,
        }
    }

    pub fn add(&mut self, task: Task) {
        info!("Adding task: {:?}", &task);
        self.task_list.push_back(task);
    }

    pub fn idle(&self) -> bool {
        self.handle.is_none()
    }

    pub fn poll(&mut self) -> Result<()> {
        // Forward any actions to main app
        if let Ok(action) = self.task_rx.try_recv() {
            let result = self.action_tx.send(action.clone());
            assert!(result.is_ok(), "Failed to send Action: {action:?}");
        }

        // Check status of current task (if one is running).
        // NOTE: Action::NextScreen is sent once all tasks are complete
        if let Some(handle) = self.handle.take() {
            if handle.is_finished() {
                if self.task_list.is_empty() {
                    // No tasks remain
                    self.task_tx.send(Action::NextScreen)?;
                } else {
                    // Start next task
                    self.start()?;
                }
            } else {
                // Task not complete, return handle
                self.handle = Some(handle);
            }
        } else if !self.task_list.is_empty() {
            // No current task but one is available
            self.start()?;
        }
        Ok(())
    }

    pub fn start(&mut self) -> Result<()> {
        if let Some(task) = self.task_list.pop_front() {
            let task_str = format!("{task:?}");
            let task_tx = self.task_tx.clone();
            match task {
                Task::Command(ref cmd_path, ref cmd_args) => {
                    let cmd_path = cmd_path.clone();
                    let cmd_args = cmd_args.clone();
                    if cfg!(windows) {
                        self.handle = Some(thread::spawn(move || {
                            let result = Command::new(cmd_path)
                                .args(cmd_args)
                                .stdout(Stdio::piped())
                                .output();
                            if let Some(action) = match result {
                                Ok(output) => {
                                    if output.status.success() {
                                        None
                                    } else {
                                        // Command returned an error status
                                        let mut msg = String::new();
                                        if let Ok(stdout) = String::from_utf8(output.stdout) {
                                            msg = String::from(stdout.trim());
                                        }
                                        if msg.is_empty() {
                                            msg = String::from("Generic error");
                                        }
                                        Some(Action::Error(format!("Command failed: {msg}",)))
                                    }
                                }
                                Err(err) => {
                                    Some(Action::Error(format!("Failed to run command: {err:?}")))
                                }
                            } {
                                let msg = format!("{:?}", &action);
                                let result = task_tx.send(action);
                                assert!(result.is_ok(), "Failed to send Action: {msg}");
                            }
                        }));
                    } else {
                        // Simulate task if not running under Windows
                        self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250))));
                    }
                }
                Task::Diskpart(ref script) => {
                    if cfg!(windows) {
                        let script = String::from(script);
                        self.handle = Some(thread::spawn(move || {
                            let output = diskpart::run_script_raw(script.as_str());
                            if !output.status.success()
                                && task_tx
                                    .send(Action::Error(String::from(
                                        "Diskpart script returned an error",
                                    )))
                                    .is_err()
                            {
                                panic!("Failed to send Action: {task_str:?}");
                            }
                        }));
                    } else {
                        // Simulate task if not running under Windows
                        self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250))));
                    }
                }
                Task::ScanDisks => {
                    let disk_list_arc = self.disk_list.clone();

                    // Queue UpdateDiskList for various components
                    self.add(Task::UpdateDiskList);

                    self.handle = Some(thread::spawn(move || {
                        let mut disks = disk_list_arc.lock().unwrap();
                        *disks = disk::get_disks();
                    }));
                }
                Task::Sleep => {
                    self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250))));
                }
                Task::UpdateDestDisk(index) => {
                    self.action_tx.send(Action::DisplayPopup(
                        popup::Type::Info,
                        String::from("Refreshing disk info"),
                    ))?;

                    // Queue UpdateDiskList for various components
                    self.add(Task::Sleep);
                    self.add(Task::UpdateDiskList);

                    // Update destination disk ~in-place
                    let disk_list_arc = self.disk_list.clone();
                    self.handle = Some(thread::spawn(move || {
                        let mut disks = disk_list_arc.lock().unwrap();
                        let old_disk = &mut disks[index];
                        disks[index] = disk::refresh_disk_info(old_disk);
                    }));
                }
                Task::UpdateDiskList => {
                    let disks = self.disk_list.lock().unwrap();
                    let disks_copy = disks.clone();
                    let action_tx = self.action_tx.clone();
                    self.handle = Some(thread::spawn(move || {
                        if let Err(err) = action_tx.send(Action::UpdateDiskList(disks_copy)) {
                            panic!("Failed to send Action: {err:?}");
                        }
                    }));
                }
            }
        }
        Ok(())
    }
}