gpu-histop 0.1.0

use std::time::{Duration, Instant};

use anyhow::Result;
use crossbeam_channel::Receiver;
use eframe::egui;
use eframe::egui::{Color32, RichText, Stroke};
use egui_plot::{Legend, Line, Plot, PlotPoints};

use crate::history::History;
use crate::model::{GpuInfo, GpuProcess, GpuProcessKind, GpuSample, MetricKind};
use crate::sampler::{SamplerEvent, SamplerHandle};

#[derive(Debug, Clone, Copy)]
pub struct GuiConfig {
    pub sample_interval: Duration,
    pub frame_interval: Duration,
    pub history_retention: Duration,
    pub initial_window: Duration,
}

pub struct GuiApp {
    devices: Vec<GpuInfo>,
    backend_label: String,
    rx: Receiver<SamplerEvent>,
    sampler: Option<SamplerHandle>,
    histories: Vec<History>,
    config: GuiConfig,
    started: Instant,
    last_error: Option<String>,
    samples_seen: u64,
    selected_gpu: usize,
    window_index: usize,
    auto_follow: bool,
}

impl GuiApp {
    pub fn new(
        devices: Vec<GpuInfo>,
        backend_label: String,
        rx: Receiver<SamplerEvent>,
        sampler: SamplerHandle,
        config: GuiConfig,
    ) -> Self {
        let histories = devices
            .iter()
            .map(|_| History::new(config.history_retention))
            .collect();

        Self {
            devices,
            backend_label,
            rx,
            sampler: Some(sampler),
            histories,
            config,
            started: Instant::now(),
            last_error: None,
            samples_seen: 0,
            selected_gpu: 0,
            window_index: closest_window(config.initial_window),
            auto_follow: true,
        }
    }

    fn ingest(&mut self) {
        while let Ok(event) = self.rx.try_recv() {
            match event {
                SamplerEvent::Samples(samples) => {
                    self.samples_seen += samples.len() as u64;
                    for sample in samples {
                        if let Some(history) = self.histories.get_mut(sample.gpu_id) {
                            history.push(sample);
                        }
                    }
                }
                SamplerEvent::Error(error) => {
                    self.last_error = Some(error);
                }
            }
        }
        if self.selected_gpu >= self.devices.len() {
            self.selected_gpu = self.devices.len().saturating_sub(1);
        }
    }

    fn selected_history(&self) -> Option<&History> {
        self.histories.get(self.selected_gpu)
    }

    fn visible_window(&self) -> Duration {
        GUI_WINDOWS[self.window_index]
    }

    fn total_stored_samples(&self) -> usize {
        self.histories.iter().map(History::len).sum()
    }

    fn draw_top_bar(&mut self, ui: &mut egui::Ui) {
        ui.horizontal_wrapped(|ui| {
            ui.heading("gpu-histop");
            ui.separator();
            ui.label(format!("{} GPU(s)", self.devices.len()));
            ui.label(format!("backend {}", self.backend_label));
            ui.label(format!("sample {:.1} Hz", hz(self.config.sample_interval)));
            ui.label(format!("stored {}", self.total_stored_samples()));
            ui.label(format!("up {}", format_duration(self.started.elapsed())));
            ui.separator();
            ui.checkbox(&mut self.auto_follow, "follow");

            egui::ComboBox::from_id_salt("history_window")
                .selected_text(format_duration(self.visible_window()))
                .show_ui(ui, |ui| {
                    for (index, window) in GUI_WINDOWS.iter().enumerate() {
                        ui.selectable_value(
                            &mut self.window_index,
                            index,
                            format_duration(*window),
                        );
                    }
                });
        });

        if let Some(error) = &self.last_error {
            ui.colored_label(Color32::from_rgb(230, 80, 70), error);
        }
    }

    fn draw_sidebar(&mut self, ui: &mut egui::Ui) {
        ui.heading("GPUs");
        ui.add_space(6.0);

        for (index, device) in self.devices.iter().enumerate() {
            let latest = self.histories.get(index).and_then(History::latest);
            let title = format!("GPU {}  {}", device.id, device.name);
            let summary = latest
                .map(gpu_sidebar_summary)
                .unwrap_or_else(|| "waiting for samples".to_owned());

            let response = ui.selectable_label(self.selected_gpu == index, title);
            ui.label(RichText::new(summary).small().color(Color32::GRAY));
            ui.add_space(8.0);
            if response.clicked() {
                self.selected_gpu = index;
            }
        }
    }

    fn draw_main(&mut self, ui: &mut egui::Ui) {
        let Some(device) = self.devices.get(self.selected_gpu) else {
            ui.centered_and_justified(|ui| {
                ui.label("No GPUs found");
            });
            return;
        };
        let Some(history) = self.selected_history() else {
            ui.label("No history available");
            return;
        };

        ui.horizontal_wrapped(|ui| {
            ui.heading(format!("GPU {} {}", device.id, device.name));
            if let Some(uuid) = &device.uuid {
                ui.label(RichText::new(short_uuid(uuid)).color(Color32::GRAY));
            }
        });

        if let Some(sample) = history.latest() {
            draw_metric_strip(ui, sample);
        }

        ui.separator();
        egui::ScrollArea::vertical()
            .auto_shrink([false, false])
            .show(ui, |ui| {
                self.draw_metric_plots(ui, history);
                ui.separator();
                self.draw_process_table(ui, history.latest());
            });
    }

    fn draw_metric_plots(&self, ui: &mut egui::Ui, history: &History) {
        let columns = if ui.available_width() >= 980.0 { 2 } else { 1 };
        ui.columns(columns, |cols| {
            for (index, metric) in MetricKind::ALL.iter().copied().enumerate() {
                let column = index % columns;
                draw_metric_plot(
                    &mut cols[column],
                    history,
                    metric,
                    self.selected_gpu,
                    self.visible_window(),
                );
            }
        });
    }

    fn draw_process_table(&self, ui: &mut egui::Ui, sample: Option<&GpuSample>) {
        ui.heading("Processes");
        ui.label(
            RichText::new("NVML compute, graphics, and MPS contexts; sorted by GPU memory.")
                .small()
                .color(Color32::GRAY),
        );

        let Some(sample) = sample else {
            ui.label("Waiting for samples.");
            return;
        };

        if sample.processes.is_empty() {
            ui.label("No processes reported for this GPU.");
            return;
        }

        egui::ScrollArea::vertical()
            .max_height(260.0)
            .auto_shrink([false, false])
            .show(ui, |ui| {
                egui::Grid::new("process_table")
                    .striped(true)
                    .num_columns(6)
                    .spacing([16.0, 6.0])
                    .show(ui, |ui| {
                        table_header(ui, "TYPE");
                        table_header(ui, "PID");
                        table_header(ui, "USER");
                        table_header(ui, "GPU MEM");
                        table_header(ui, "MIG");
                        table_header(ui, "COMMAND");
                        ui.end_row();

                        for process in &sample.processes {
                            ui.colored_label(process_color(process), process.kind_label());
                            ui.monospace(process.pid.to_string());
                            ui.label(process.user.as_deref().unwrap_or("?"));
                            ui.label(
                                process
                                    .used_gpu_memory_bytes
                                    .map(format_bytes_compact)
                                    .unwrap_or_else(|| "n/a".to_owned()),
                            );
                            ui.label(process_mig_label(process));
                            ui.monospace(process.command.as_deref().unwrap_or("?"));
                            ui.end_row();
                        }
                    });
            });
    }
}

impl Drop for GuiApp {
    fn drop(&mut self) {
        if let Some(sampler) = self.sampler.take() {
            sampler.stop();
        }
    }
}

impl eframe::App for GuiApp {
    fn ui(&mut self, ui: &mut egui::Ui, _frame: &mut eframe::Frame) {
        self.ingest();

        egui::Panel::top("top_bar")
            .resizable(false)
            .show_inside(ui, |ui| {
                ui.add_space(6.0);
                self.draw_top_bar(ui);
                ui.add_space(6.0);
            });

        egui::Panel::left("gpu_sidebar")
            .resizable(true)
            .default_size(260.0)
            .size_range(220.0..=380.0)
            .show_inside(ui, |ui| {
                self.draw_sidebar(ui);
            });

        egui::CentralPanel::default().show_inside(ui, |ui| {
            self.draw_main(ui);
        });

        if self.auto_follow {
            ui.ctx().request_repaint_after(self.config.frame_interval);
        }
    }
}

pub fn run(app: GuiApp) -> Result<()> {
    let native_options = eframe::NativeOptions {
        viewport: egui::ViewportBuilder::default()
            .with_title("gpu-histop")
            .with_inner_size([1420.0, 920.0])
            .with_min_inner_size([980.0, 660.0]),
        ..Default::default()
    };

    eframe::run_native(
        "gpu-histop",
        native_options,
        Box::new(move |_cc| Ok(Box::new(app))),
    )?;
    Ok(())
}

fn draw_metric_strip(ui: &mut egui::Ui, sample: &GpuSample) {
    ui.horizontal_wrapped(|ui| {
        metric_badge(
            ui,
            "GPU",
            sample.gpu_util_percent,
            "%",
            metric_color(MetricKind::GpuUtil),
        );
        metric_badge(
            ui,
            "MEM",
            sample.mem_util_percent,
            "%",
            metric_color(MetricKind::MemUtil),
        );
        ui.label(format!(
            "VRAM {}",
            vram_summary(sample.vram_used_bytes, sample.vram_total_bytes)
        ));
        metric_badge(
            ui,
            "PWR",
            sample.power_watts,
            "W",
            metric_color(MetricKind::Power),
        );
        if let Some(limit) = sample.power_limit_watts {
            ui.label(RichText::new(format!("/{limit:.0}W")).color(Color32::GRAY));
        }
        metric_badge(
            ui,
            "TEMP",
            sample.temperature_celsius,
            "C",
            metric_color(MetricKind::Temperature),
        );
        metric_badge(
            ui,
            "FAN",
            sample.fan_percent,
            "%",
            metric_color(MetricKind::Fan),
        );
        ui.label(format!(
            "CLK {}/{} MHz",
            whole_or_na(sample.graphics_clock_mhz),
            whole_or_na(sample.memory_clock_mhz)
        ));
        ui.label(format!("PROC {}", sample.processes.len()));
    });
}

fn metric_badge(ui: &mut egui::Ui, label: &str, value: Option<f64>, unit: &str, color: Color32) {
    let value = value
        .map(|value| format!("{value:.0}{unit}"))
        .unwrap_or_else(|| "n/a".to_owned());
    ui.label(
        RichText::new(format!("{label} {value}"))
            .color(color)
            .strong(),
    );
}

fn draw_metric_plot(
    ui: &mut egui::Ui,
    history: &History,
    metric: MetricKind,
    gpu_id: usize,
    window: Duration,
) {
    let now = Instant::now();
    let latest = history.latest();
    let points = metric_points(history, metric, now, window);
    let latest_value = latest.and_then(|sample| metric.value(sample));
    let title = latest_value
        .map(|value| format!("{}  {:.1}{}", metric.title(), value, metric.unit()))
        .unwrap_or_else(|| format!("{}  n/a", metric.title()));
    let color = metric_color(metric);
    let (min_y, max_y) = metric_scale(history, metric, latest, now, window);

    ui.group(|ui| {
        ui.horizontal(|ui| {
            ui.label(RichText::new(title).strong().color(color));
            ui.label(
                RichText::new(format_duration(window))
                    .small()
                    .color(Color32::GRAY),
            );
        });

        Plot::new(format!("gpu_{gpu_id}_{metric:?}"))
            .height(170.0)
            .legend(Legend::default())
            .show_grid(true)
            .allow_drag(true)
            .allow_zoom(true)
            .include_x(-(window.as_secs_f64()))
            .include_x(0.0)
            .include_y(min_y)
            .include_y(max_y)
            .label_formatter(move |name, point| {
                if name.is_empty() {
                    format!("{:.2}s, {:.2}{}", point.x, point.y, metric.unit())
                } else {
                    format!("{name}\n{:.2}s, {:.2}{}", point.x, point.y, metric.unit())
                }
            })
            .show(ui, |plot_ui| {
                if !points.is_empty() {
                    plot_ui.line(
                        Line::new(metric.title(), points)
                            .color(color)
                            .stroke(Stroke::new(1.6, color)),
                    );
                }
            });
    });
}

fn metric_points(
    history: &History,
    metric: MetricKind,
    now: Instant,
    window: Duration,
) -> PlotPoints<'static> {
    history
        .iter_window(now, window)
        .filter_map(|sample| {
            let value = metric.value(sample)?;
            let age = if sample.at <= now {
                now.duration_since(sample.at).as_secs_f64()
            } else {
                0.0
            };
            Some([-age, value])
        })
        .collect::<Vec<_>>()
        .into()
}

fn metric_scale(
    history: &History,
    metric: MetricKind,
    latest: Option<&GpuSample>,
    now: Instant,
    window: Duration,
) -> (f64, f64) {
    if let Some(range) = metric.fixed_range(latest) {
        return range;
    }

    let mut max_seen: f64 = 0.0;
    for sample in history.iter_window(now, window) {
        if let Some(value) = metric.value(sample) {
            max_seen = max_seen.max(value);
        }
    }
    (0.0, nice_ceiling((max_seen * 1.15).max(1.0)))
}

fn table_header(ui: &mut egui::Ui, text: &str) {
    ui.label(RichText::new(text).small().strong().color(Color32::GRAY));
}

fn gpu_sidebar_summary(sample: &GpuSample) -> String {
    format!(
        "GPU {} | VRAM {} | {} proc | PWR {} | TEMP {}",
        percent_or_na(sample.gpu_util_percent),
        sample
            .vram_used_percent()
            .map(|v| format!("{v:.0}%"))
            .unwrap_or_else(|| "n/a".to_owned()),
        sample.processes.len(),
        sample
            .power_watts
            .map(|v| format!("{v:.0}W"))
            .unwrap_or_else(|| "n/a".to_owned()),
        sample
            .temperature_celsius
            .map(|v| format!("{v:.0}C"))
            .unwrap_or_else(|| "n/a".to_owned())
    )
}

fn metric_color(metric: MetricKind) -> Color32 {
    match metric {
        MetricKind::GpuUtil => Color32::from_rgb(65, 200, 230),
        MetricKind::MemUtil => Color32::from_rgb(80, 210, 140),
        MetricKind::VramUsed => Color32::from_rgb(230, 190, 75),
        MetricKind::Power => Color32::from_rgb(210, 120, 230),
        MetricKind::Temperature => Color32::from_rgb(235, 100, 85),
        MetricKind::Fan => Color32::from_rgb(115, 150, 245),
    }
}

fn process_color(process: &GpuProcess) -> Color32 {
    if process.kinds.contains(&GpuProcessKind::Graphics) {
        Color32::from_rgb(230, 190, 75)
    } else if process.kinds.contains(&GpuProcessKind::Mps) {
        Color32::from_rgb(210, 120, 230)
    } else {
        Color32::from_rgb(80, 210, 140)
    }
}

fn process_mig_label(process: &GpuProcess) -> String {
    match (process.gpu_instance_id, process.compute_instance_id) {
        (Some(gpu), Some(compute)) => format!("gi={gpu}/ci={compute}"),
        (Some(gpu), None) => format!("gi={gpu}"),
        _ => String::new(),
    }
}

fn vram_summary(used: Option<u64>, total: Option<u64>) -> String {
    match (used, total) {
        (Some(used), Some(total)) if total > 0 => format!(
            "{:.1}/{:.1} GiB {:.0}%",
            bytes_to_gib(used),
            bytes_to_gib(total),
            used as f64 * 100.0 / total as f64
        ),
        (Some(used), _) => format!("{:.1} GiB used", bytes_to_gib(used)),
        _ => "n/a".to_owned(),
    }
}

fn bytes_to_gib(bytes: u64) -> f64 {
    bytes as f64 / 1024.0 / 1024.0 / 1024.0
}

fn format_bytes_compact(bytes: u64) -> String {
    let gib = bytes_to_gib(bytes);
    if gib >= 10.0 {
        format!("{gib:.0} GiB")
    } else if gib >= 1.0 {
        format!("{gib:.1} GiB")
    } else {
        format!("{:.0} MiB", bytes as f64 / 1024.0 / 1024.0)
    }
}

fn percent_or_na(value: Option<f64>) -> String {
    value
        .map(|value| format!("{value:.0}%"))
        .unwrap_or_else(|| "n/a".to_owned())
}

fn whole_or_na(value: Option<f64>) -> String {
    value
        .map(|v| format!("{v:.0}"))
        .unwrap_or_else(|| "n/a".to_owned())
}

fn short_uuid(uuid: &str) -> &str {
    uuid.rsplit('-').next().unwrap_or(uuid)
}

fn hz(duration: Duration) -> f64 {
    1.0 / duration.as_secs_f64().max(0.001)
}

fn format_duration(duration: Duration) -> String {
    let seconds = duration.as_secs();
    if seconds < 60 {
        format!("{seconds}s")
    } else if seconds < 3600 {
        format!("{}m{:02}s", seconds / 60, seconds % 60)
    } else {
        format!("{}h{:02}m", seconds / 3600, (seconds % 3600) / 60)
    }
}

fn nice_ceiling(value: f64) -> f64 {
    if value <= 10.0 {
        return 10.0;
    }

    let magnitude = 10_f64.powf(value.log10().floor());
    let normalized = value / magnitude;
    let rounded = if normalized <= 2.0 {
        2.0
    } else if normalized <= 5.0 {
        5.0
    } else {
        10.0
    };
    rounded * magnitude
}

const GUI_WINDOWS: [Duration; 8] = [
    Duration::from_secs(10),
    Duration::from_secs(30),
    Duration::from_secs(60),
    Duration::from_secs(5 * 60),
    Duration::from_secs(10 * 60),
    Duration::from_secs(30 * 60),
    Duration::from_secs(60 * 60),
    Duration::from_secs(3 * 60 * 60),
];

fn closest_window(target: Duration) -> usize {
    GUI_WINDOWS
        .iter()
        .enumerate()
        .min_by_key(|(_, window)| window.abs_diff(target))
        .map(|(index, _)| index)
        .unwrap_or(2)
}