nix-output/src/handlers/substitute_status.rs

use std::{borrow::Cow, collections::HashMap, time::Instant};

use indicatif::{HumanBytes, ProgressBar};
use owo_colors::Style;

use crate::{
    DONE_CHAR, DOWNLOAD_CHAR, INPROGRESS_CHAR,
    action::{Action, ActionType, BuildStepId, ResultFields, StartFields},
    download_display::{DOWNLOAD_STYLE, EXTRACT_STYLE, StyledString, style_bar},
    estimator::Estimator,
    handlers::Handler,
    multibar::{BarSegment, MultiBar},
};

pub struct CopyPathsHandler;

impl Handler for CopyPathsHandler {
    fn handle(
        &mut self,
        state: &mut crate::state::State,
        action: &crate::action::Action,
    ) -> color_eyre::Result<bool> {
        match action {
            Action::Start {
                start_type: StartFields::CopyPaths,
                id,
                ..
            } => {
                let progress = state.add_pb(ProgressBar::new(1));

                progress.set_style(indicatif::ProgressStyle::with_template("{msg}").unwrap());
                state.plug(SubstitutionStatusHandler::new(
                    id,
                    state.term_width,
                    progress,
                ));
                // Handle global start action
                Ok(true)
            }
            _ => Ok(true),
        }
    }
}

pub struct SubstitutionStatusHandler {
    id: BuildStepId,
    progress: ProgressBar,
    download_estimator: Estimator,
    extract_estimator: Estimator,
    state_copy: HashMap<BuildStepId, [u64; 2]>,
    state_transfer: HashMap<BuildStepId, [u64; 2]>,
    state_self: [u64; 2],
    max_copy: u64,
    max_transfer: u64,
}

impl SubstitutionStatusHandler {
    fn new(id: &BuildStepId, width: u16, progress: ProgressBar) -> Self {
        let mut new = Self {
            id: *id,
            progress,
            download_estimator: Estimator::new(Instant::now()),
            extract_estimator: Estimator::new(Instant::now()),
            state_copy: HashMap::new(),
            state_transfer: HashMap::new(),
            state_self: [0, 0],
            max_copy: 0,
            max_transfer: 0,
        };
        new.draw_bar(width);
        new
    }

    fn get_done(&self) -> u64 {
        self.state_transfer.values().map(|&[done, ..]| done).sum()
    }

    fn get_running(&self) -> u64 {
        self.state_transfer
            .values()
            .map(|&[_, expected, ..]| expected)
            .sum()
    }
    fn get_running_copy(&self) -> u64 {
        self.state_copy
            .values()
            .map(|&[_, expected, ..]| expected)
            .sum()
    }

    fn get_unpacked(&self) -> u64 {
        self.state_copy.values().map(|&[done, ..]| done).sum()
    }

    fn draw_bar(&mut self, width: u16) {
        if self.max_transfer == 0 || self.max_copy == 0 {
            return;
        }
        let dl_percent = ((self.get_done() as f64 / self.max_transfer as f64) * 100.0) as u64;
        let ex_percent = ((self.get_unpacked() as f64 / self.max_copy as f64) * 100.0) as u64;

        let dl_expected_percent =
            ((self.get_running() as f64 / self.max_transfer as f64) * 100.0) as u64;
        let ex_expected_percent =
            ((self.get_running_copy() as f64 / self.max_copy as f64) * 100.0) as u64;

        let min = dl_percent.min(ex_percent);
        let dl = dl_percent.saturating_sub(min);
        let expected = dl_expected_percent.max(ex_expected_percent);
        let exp = expected.saturating_sub(min + dl);
        let mbar = MultiBar([
            BarSegment::Dynamic(&DONE_CHAR, min),
            BarSegment::Dynamic(&DOWNLOAD_CHAR, dl),
            BarSegment::Dynamic(&INPROGRESS_CHAR, exp),
            BarSegment::Dynamic(" ", 100 - min - dl - exp),
        ]);

        let work_per_sec = if self.get_done() < self.max_transfer {
            StyledString::new(
                Cow::Owned(format!(
                    "{}",
                    HumanBytes(self.download_estimator.steps_per_second(Instant::now()) as u64)
                )),
                DOWNLOAD_STYLE,
            )
        } else {
            StyledString::new(
                Cow::Owned(format!(
                    "{}",
                    HumanBytes(self.extract_estimator.steps_per_second(Instant::now()) as u64)
                )),
                EXTRACT_STYLE,
            )
        };
        let msg = style_bar(
            StyledString::new(Cow::Owned("Downloading".to_string()), Style::new()),
            StyledString::new(
                Cow::Owned(format!(
                    "{}/{} {min} {dl} {expected}",
                    self.state_self[0], self.state_self[1]
                )),
                Style::new().purple(),
            ),
            work_per_sec,
            self.get_done() as usize,
            self.get_unpacked() as usize,
            mbar,
            width,
        );
        self.progress.set_message(msg);
        self.progress.tick();
    }
}

impl Handler for SubstitutionStatusHandler {
    fn handle(
        &mut self,
        state: &mut crate::state::State,
        action: &crate::action::Action,
    ) -> color_eyre::Result<bool> {
        match action {
            Action::Start {
                start_type: StartFields::CopyPath { path, .. },
                id,
                ..
            } => {
                self.state_copy.insert(*id, [0; 2]);
                Ok(true)
            }
            Action::Start {
                start_type: StartFields::FileTransfer { .. },
                id,
                ..
            } => {
                self.state_transfer.insert(*id, [0; 2]);
                Ok(true)
            }

            Action::Result {
                id,
                fields:
                    ResultFields::SetExpected {
                        action: ActionType::FileTransfer,
                        expected,
                    },
            } => {
                self.max_transfer = *expected;
                self.draw_bar(state.term_width);
                Ok(true)
            }
            Action::Result {
                id,
                fields:
                    ResultFields::SetExpected {
                        action: ActionType::CopyPath,
                        expected,
                    },
            } => {
                self.max_copy = *expected;
                self.draw_bar(state.term_width);
                Ok(true)
            }

            Action::Result {
                id,
                fields: ResultFields::Progress { done, expected, .. },
            } => {
                if id == &self.id {
                    self.state_self = [*done, *expected];
                    self.draw_bar(state.term_width);
                }
                if let Some(copy) = self.state_copy.get_mut(id) {
                    *copy = [*done, *expected];
                    self.extract_estimator
                        .record(self.get_unpacked(), Instant::now());
                    self.draw_bar(state.term_width);
                }

                if let Some(transfer) = self.state_transfer.get_mut(id) {
                    *transfer = [*done, *expected];

                    self.download_estimator
                        .record(self.get_done(), Instant::now());
                    self.draw_bar(state.term_width);
                }
                Ok(true)
            }
            _ => Ok(true),
        }
    }
}