nix-output/src/handlers/download.rs

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

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

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

pub struct SubstituteHandler;

impl Handler for SubstituteHandler {
    fn handle(
        &mut self,
        state: &mut crate::state::State,
        action: &crate::action::Action,
    ) -> color_eyre::Result<bool> {
        match action {
            Action::Start {
                start_type: StartFields::Substitute { source, target },
                id,
                ..
            } => {
                state.plug(CopyHandler(*id));
                // Handle global start action
                Ok(true)
            }
            _ => Ok(true),
        }
    }
}

pub struct CopyHandler(BuildStepId);
impl Handler for CopyHandler {
    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,
                        origin,
                        destination,
                    },
                id,
                parent,
                ..
            } if parent == &self.0 => {
                state.plug(DownloadHandler::new(
                    state.term_width,
                    *id,
                    path.to_string(),
                ));
                // Handle global start action
                Ok(false)
            }
            _ => Ok(true),
        }
    }
}

pub struct DownloadHandler {
    pub id: BuildStepId,
    pub copy_id: BuildStepId,
    pub download_estimator: Estimator,
    pub extract_estimator: Estimator,
    pub bar: Option<ProgressBar>,
    pub path: String,
    pub download_expected: u64,
    pub download_done: u64,
    pub extract_expected: u64,
    pub extract_done: u64,
    pub started_at: Instant,
}

impl DownloadHandler {
    pub fn new(width: u16, copy_id: BuildStepId, path: String) -> Self {
        let new = Self {
            id: u64::MAX.into(),
            copy_id,
            download_estimator: Estimator::new(Instant::now()),
            extract_estimator: Estimator::new(Instant::now()),
            bar: None,
            path,
            download_expected: 0,
            download_done: 0,
            extract_expected: 0,
            extract_done: 0,
            started_at: Instant::now(),
        };
        new.draw_bar(width);
        new
    }
    pub fn draw_bar(&self, width: u16) {
        let Some(bar) = &self.bar else {
            return;
        };
        let name = nix_path::extract_package_name_string(&self.path)
            .unwrap_or_else(|| "unknown".to_string());

        if self.download_expected == 0 || self.extract_expected == 0 {
            bar.set_message(format!("Downloading {name}"));
            bar.tick();
            return;
        }

        let work_per_sec = if self.download_done < self.download_expected {
            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 dl_percent =
            ((self.download_done as f64 / self.download_expected as f64) * 100.0) as u64;
        let ex_percent = ((self.extract_done as f64 / self.extract_expected as f64) * 100.0) as u64;
        let min = dl_percent.min(ex_percent);
        let dl = dl_percent.saturating_sub(min);
        let mbar = MultiBar([
            BarSegment::Dynamic(&DONE_CHAR, min),
            BarSegment::Dynamic(&DOWNLOAD_CHAR, dl),
            BarSegment::Dynamic(" ", 100 - min - dl),
        ]);
        let msg = style_bar(
            StyledString::new(Cow::Owned("Downloading".to_string()), Style::new()),
            StyledString::new(Cow::Owned(name), Style::new().purple()),
            work_per_sec,
            self.download_done as usize,
            self.extract_done as usize,
            mbar,
            width,
        );

        // let total_expected = self.download_expected + self.extract_expected;
        // let total_done = self.download_done + self.extract_done;
        //
        //
        // let download_per_sec =
        //     HumanBytes(self.download_estimator.steps_per_second(Instant::now()) as u64).to_string();
        // let extract_per_sec =
        //     HumanBytes(self.extract_estimator.steps_per_second(Instant::now()) as u64).to_string();
        // let download_done_human = HumanBytes(self.download_done).to_string();
        // let download_done_human = pad_string(&download_done_human, INFO_WIDTH)
        //     .blue()
        //     .bold()
        //     .to_string();
        // let extract_done_human = HumanBytes(self.extract_done).to_string();
        // let extract_done_human = pad_string(&extract_done_human, INFO_WIDTH)
        //     .green()
        //     .bold()
        //     .to_string();
        // let download_per_sec = pad_string(&format!("{download_per_sec}/s"), INFO_WIDTH)
        //     .blue()
        //     .bold()
        //     .to_string();
        // let extract_per_sec = pad_string(&format!("{extract_per_sec}/s"), INFO_WIDTH)
        //     .green()
        //     .bold()
        //     .to_string();
        //
        // let display = format!("{work_per_sec} | {download_done_human} | {extract_done_human} ");
        // let display_length = (INFO_WIDTH * 3) + 9;
        //
        // // + 6 to account for final format
        // let total_length = status_len + display_length + 4;
        //
        // if width <= len {
        //     self.bar.set_message(format!(
        //         "{}: {}/{}",
        //         name,
        //         total_done,
        //         if total_expected == 0 {
        //             "-".to_string()
        //         } else {
        //             total_expected.to_string()
        //         }
        //     ));
        //     self.bar.tick();
        //     return;
        // }
        // let msg = match width {
        //     0..60 => {
        //         format!(
        //             "{}: {}/{}",
        //             name,
        //             total_done,
        //             if total_expected == 0 {
        //                 "-".to_string()
        //             } else {
        //                 total_expected.to_string()
        //             }
        //         )
        //     }
        //     _ => {
        //         .scale((width.saturating_sub(total_length as u16)) as u64);
        //         let bar = if width > display_length as u16 + status_len as u16 {
        //             format!("[{}]", bar)
        //         } else {
        //             String::new()
        //         };
        //
        //         format!("{status} {display} {bar}",)
        //     }
        // };
        bar.set_message(msg);
        bar.tick();
    }
}

impl Handler for DownloadHandler {
    fn handle(
        &mut self,
        state: &mut crate::state::State,
        action: &crate::action::Action,
    ) -> color_eyre::Result<bool> {
        match action {
            Action::Start {
                start_type: StartFields::FileTransfer { target },
                id,
                parent,
                ..
            } if parent == &self.copy_id => {
                self.id = *id;
                Ok(true)
            }
            Action::Result {
                fields:
                    ResultFields::Progress {
                        done,
                        expected,
                        running: _,
                        failed: _,
                    },
                id,
            } => {
                if id == &self.id {
                    self.download_done = *done;
                    self.download_expected = *expected;
                    self.download_estimator.record(*done, Instant::now());
                    self.draw_bar(state.term_width);
                }
                if id == &self.copy_id {
                    self.extract_done = *done;
                    self.extract_expected = *expected;
                    self.extract_estimator.record(*done, Instant::now());
                    self.draw_bar(state.term_width);
                }
                if self.download_expected > 0 && self.extract_expected > 0 && self.bar.is_none() {
                    let bar = state.add_pb(ProgressBar::new(1));
                    bar.set_style(indicatif::ProgressStyle::with_template("{msg}").unwrap());
                    self.bar = Some(bar);
                    self.draw_bar(state.term_width);
                }

                Ok(true)
            }
            Action::Stop { id } if id == &self.copy_id => {
                if let Some(bar) = &self.bar {
                    bar.finish_and_clear();
                }
                Ok(false)
            }
            _ => Ok(true),
        }
    }
    fn on_resize(&mut self, state: &mut crate::state::State) -> color_eyre::Result<()> {
        self.draw_bar(state.term_width);
        Ok(())
    }
}