bandwidthmon 0.1.18

//! Bandwidth Monitor 3 - Cross-platform using sysinfo + rasciichart
//! Author: Hadi Cahyadi <cumulus13@gmail.com>
//! License: MIT

use anyhow::{Context, Result};
use clap::Parser;
use clap::ArgAction;
use crossterm::{
    cursor::{Hide, MoveTo, Show},
    event::{self, Event, KeyCode},
    execute,
    style::{Color, Print},
    terminal::{
        disable_raw_mode, enable_raw_mode, size, EnterAlternateScreen,
        LeaveAlternateScreen,
    },
};
use rasciichart::{plot_with_config, Config};
use std::collections::VecDeque;
use std::io::{stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use sysinfo::Networks;
use std::fmt;

const INTERVAL: Duration = Duration::from_secs(1);
const DEFAULT_HISTORY: usize = 120;
const DEFAULT_HEIGHT: usize = 10;

struct ColoredVersion;

impl ColoredVersion {
    pub fn new() -> Self {
        Self {}
    }
}

impl fmt::Display for ColoredVersion {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = style_text("bandwidthmon3", Color::Yellow, true);
        let author = style_text("Hadi Cahyadi <cumulus13@gmail.com>", Color::Cyan, true);
        let version = style_text(env!("CARGO_PKG_VERSION"), Color::White, true);

        write!(f, "{} {} by {}", name, version, author)
    }
}

#[derive(Parser, Debug)]
#[command(
    // name = "bandwidthmon3",
    // disable_version_flag = true,
    // author = "Hadi Cahyadi <cumulus13@gmail.com>",
    // about = "Cross-platform bandwidth monitor with rasciichart",
    // long_about = None,
    // version = ColoredVersion::new().as_str()
    about = "Cross-platform bandwidth monitor with rasciichart",
    disable_version_flag = true
)]
struct Args {
    /// Network interface to monitor (auto-select if not specified)
    #[arg(short, long)]
    iface: Option<String>,

    /// Chart height in lines
    #[arg(short = 'H', long, default_value_t = DEFAULT_HEIGHT)]
    height: usize,

    /// Chart width in columns (auto-fit terminal if 0)
    #[arg(short = 'W', long, default_value_t = 0)]
    width: usize,

    /// List available network interfaces
    #[arg(short, long)]
    list: bool,

    /// Show summary statistics
    #[arg(short, long)]
    summary: bool,

    /// Show download chart only
    #[arg(short, long)]
    download: bool,

    /// Show upload chart only
    #[arg(short, long)]
    upload: bool,

    /// Maximum history points
    #[arg(long, default_value_t = DEFAULT_HISTORY)]
    history: usize,

    #[arg(short = 'v', short = 'V', long = "version", action = ArgAction::SetTrue)]
    version: bool,
}

#[derive(Debug, Clone)]
struct BandwidthStats {
    download_bps: f64,
    upload_bps: f64,
    total_rx: u64,
    total_tx: u64,
}

struct NetworkMonitor {
    interface: String,
    networks: Networks,
    history_dl: VecDeque<f64>,
    history_ul: VecDeque<f64>,
    prev_rx: u64,
    prev_tx: u64,
    prev_time: Instant,
    start_time: Instant,
    peak_dl: f64,
    peak_ul: f64,
    avg_dl: f64,
    avg_ul: f64,
    sample_count: u64,
}

impl NetworkMonitor {
    fn new(interface: String, history_size: usize) -> Result<Self> {
        let networks = Networks::new_with_refreshed_list();
        
        if !networks.iter().any(|(name, _)| name == &interface) {
            anyhow::bail!("Interface '{}' not found", interface);
        }

        let (prev_rx, prev_tx) = networks
            .get(&interface)
            .map(|data| (data.total_received(), data.total_transmitted()))
            .unwrap_or((0, 0));

        let now = Instant::now();

        Ok(Self {
            interface,
            networks,
            history_dl: VecDeque::with_capacity(history_size),
            history_ul: VecDeque::with_capacity(history_size),
            prev_rx,
            prev_tx,
            prev_time: now,
            start_time: now,
            peak_dl: 0.0,
            peak_ul: 0.0,
            avg_dl: 0.0,
            avg_ul: 0.0,
            sample_count: 0,
        })
    }

    fn update(&mut self) -> Result<BandwidthStats> {
        self.networks.refresh(false);

        let data = self
            .networks
            .get(&self.interface)
            .context("Interface disappeared")?;

        let cur_rx = data.total_received();
        let cur_tx = data.total_transmitted();
        let cur_time = Instant::now();

        let elapsed = cur_time.duration_since(self.prev_time).as_secs_f64();
        
        if elapsed < 0.001 {
            return Ok(BandwidthStats {
                download_bps: 0.0,
                upload_bps: 0.0,
                total_rx: cur_rx,
                total_tx: cur_tx,
            });
        }

        let dl_bytes = cur_rx.saturating_sub(self.prev_rx);
        let ul_bytes = cur_tx.saturating_sub(self.prev_tx);

        let dl_bps = (dl_bytes as f64) / elapsed;
        let ul_bps = (ul_bytes as f64) / elapsed;

        self.prev_rx = cur_rx;
        self.prev_tx = cur_tx;
        self.prev_time = cur_time;

        // Update history
        if self.history_dl.len() >= self.history_dl.capacity() {
            self.history_dl.pop_front();
        }
        self.history_dl.push_back(dl_bps);

        if self.history_ul.len() >= self.history_ul.capacity() {
            self.history_ul.pop_front();
        }
        self.history_ul.push_back(ul_bps);

        // Update statistics
        self.peak_dl = self.peak_dl.max(dl_bps);
        self.peak_ul = self.peak_ul.max(ul_bps);

        self.sample_count += 1;
        self.avg_dl += (dl_bps - self.avg_dl) / self.sample_count as f64;
        self.avg_ul += (ul_bps - self.avg_ul) / self.sample_count as f64;

        Ok(BandwidthStats {
            download_bps: dl_bps,
            upload_bps: ul_bps,
            total_rx: cur_rx,
            total_tx: cur_tx,
        })
    }

    fn get_history_dl(&self) -> Vec<f64> {
        self.history_dl.iter().copied().collect()
    }

    fn get_history_ul(&self) -> Vec<f64> {
        self.history_ul.iter().copied().collect()
    }
}

fn list_interfaces() -> Result<()> {
    let networks = Networks::new_with_refreshed_list();
    
    println!("\n{}", style_text("Available Network Interfaces:", Color::Cyan, true));
    println!("{}", "─".repeat(80));

    for (name, data) in networks.iter() {
        let rx = data.total_received();
        let tx = data.total_transmitted();
        let status = if rx > 0 || tx > 0 { "active" } else { "inactive" };
        
        println!(
            "  {} {} {}",
            style_text(name, Color::White, true),
            style_text(
                &format!("(RX: {}, TX: {})", 
                    format_total_bytes(rx), 
                    format_total_bytes(tx)
                ),
                Color::DarkGrey,
                false
            ),
            style_text(&format!("[{}]", status), Color::Green, false)
        );
    }
    println!();

    Ok(())
}

fn select_best_interface() -> Result<String> {
    let networks = Networks::new_with_refreshed_list();
    
    // First, try to find the most active interface
    let best = networks
        .iter()
        .filter(|(name, _)| {
            // Skip loopback interfaces
            !name.starts_with("lo") && !name.starts_with("Local")
        })
        .max_by_key(|(_, data)| data.total_received() + data.total_transmitted())
        .map(|(name, _)| name.clone());

    if let Some(interface) = best {
        return Ok(interface);
    }

    // If no active interface found, just return the first non-loopback
    networks
        .iter()
        .find(|(name, _)| !name.starts_with("lo") && !name.starts_with("Local"))
        .map(|(name, _)| name.clone())
        .context("No suitable network interfaces found")
}

fn resolve_interface(pattern: &str) -> Result<String> {
    let networks = Networks::new_with_refreshed_list();
    let interfaces: Vec<String> = networks.iter().map(|(name, _)| name.clone()).collect();
    
    // 1. Exact match
    if interfaces.iter().any(|name| name == pattern) {
        return Ok(pattern.to_string());
    }
    
    // 2. Case-insensitive partial match
    let pattern_lower = pattern.to_lowercase();
    let matches: Vec<String> = interfaces
        .iter()
        .filter(|name| name.to_lowercase().contains(&pattern_lower))
        .cloned()
        .collect();
    
    if matches.is_empty() {
        anyhow::bail!(
            "No interface matches '{}'. Available interfaces:\n  {}",
            pattern,
            interfaces.join("\n  ")
        );
    }
    
    if matches.len() == 1 {
        return Ok(matches[0].clone());
    }
    
    // Multiple matches - return the shortest one (most specific)
    Ok(matches
        .into_iter()
        .min_by_key(|s| s.len())
        .unwrap())
}

fn format_bytes(bytes: f64) -> String {
    const UNITS: &[&str] = &["B/s", "KB/s", "MB/s", "GB/s"];
    let mut value = bytes;
    let mut unit_idx = 0;

    while value >= 1024.0 && unit_idx < UNITS.len() - 1 {
        value /= 1024.0;
        unit_idx += 1;
    }

    format!("{:>7.2} {}", value, UNITS[unit_idx])
}

fn format_total_bytes(bytes: u64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut value = bytes as f64;
    let mut unit_idx = 0;

    while value >= 1024.0 && unit_idx < UNITS.len() - 1 {
        value /= 1024.0;
        unit_idx += 1;
    }

    format!("{:.2} {}", value, UNITS[unit_idx])
}

fn style_text(text: &str, color: Color, bold: bool) -> String {
    if bold {
        format!("\x1b[1m\x1b[38;5;{}m{}\x1b[0m", color_to_256(color), text)
    } else {
        format!("\x1b[38;5;{}m{}\x1b[0m", color_to_256(color), text)
    }
}

fn color_to_256(color: Color) -> u8 {
    match color {
        Color::Cyan => 51,
        Color::Yellow => 226,
        Color::White => 15,
        Color::DarkGrey => 240,
        Color::Green => 46,
        Color::Magenta => 201,
        Color::Red => 196,
        _ => 15,
    }
}

/// Render chart using custom rasciichart library
fn render_chart_rasciichart(
    data: &[f64],
    height: usize,
    width: usize,
    color: Color,
    label: &str,
) -> String {
    if data.is_empty() || height == 0 || width == 0 {
        return String::new();
    }

    // Get the last `width` points for plotting
    let start_idx = data.len().saturating_sub(width);
    let plot_data: Vec<f64> = data[start_idx..].to_vec();

    if plot_data.is_empty() {
        return String::new();
    }

    // Configure rasciichart with proper width and height
    let config = Config::default()
        .with_height(height)
        .with_width(width)
        .with_labels(true)
        .with_label_format("{:.1}".to_string());

    // Generate the chart
    let chart = match plot_with_config(&plot_data, config) {
        Ok(c) => c,
        Err(e) => return format!("Chart error: {}", e),
    };

    // Add color to the chart
    let color_code = color_to_256(color);
    let colored_chart: String = chart
        .lines()
        .map(|line| format!("\x1b[38;5;{}m{}\x1b[0m", color_code, line))
        .collect::<Vec<_>>()
        .join("\n");

    // Add label
    format!(
        "{}\n{}",
        style_text(label, color, true),
        colored_chart
    )
}

fn render_ui(
    monitor: &NetworkMonitor,
    stats: &BandwidthStats,
    args: &Args,
    term_width: u16,
) -> Result<String> {
    let mut output = String::new();
    
    // Calculate chart width based on terminal width
    // Account for label width (approximately 10 chars) and margins
    let chart_width = if args.width > 0 {
        args.width
    } else {
        // Auto-resize: terminal width minus labels and margins
        term_width.saturating_sub(20).max(30) as usize
    };

    // Header
    output.push_str(&format!(
        "{}\n",
        style_text(
            &format!("═══ Bandwidth Monitor ({}) ═══", monitor.interface),
            Color::Cyan,
            true
        )
    ));

    // Current speeds
    output.push_str(&format!(
        "{} {}  │  {} {}  {}\n",
        style_text("Download:", Color::Cyan, true),
        style_text(&format_bytes(stats.download_bps), Color::White, false),
        style_text("Upload:", Color::Yellow, true),
        style_text(&format_bytes(stats.upload_bps), Color::White, false),
        style_text("'q'/Ctrl+C=quit", Color::DarkGrey, false)
    ));

    if args.summary {
        output.push_str(&format!(
            "{} {}  │  {} {}\n",
            style_text("Peak DL:", Color::Cyan, false),
            style_text(&format_bytes(monitor.peak_dl), Color::White, false),
            style_text("Peak UL:", Color::Yellow, false),
            style_text(&format_bytes(monitor.peak_ul), Color::White, false),
        ));
        output.push_str(&format!(
            "{} {}  │  {} {}\n",
            style_text("Avg DL:", Color::Cyan, false),
            style_text(&format_bytes(monitor.avg_dl), Color::White, false),
            style_text("Avg UL:", Color::Yellow, false),
            style_text(&format_bytes(monitor.avg_ul), Color::White, false),
        ));
        output.push_str(&format!(
            "{} {}  │  {} {}\n",
            style_text("Total RX:", Color::Cyan, false),
            style_text(&format_total_bytes(stats.total_rx), Color::White, false),
            style_text("Total TX:", Color::Yellow, false),
            style_text(&format_total_bytes(stats.total_tx), Color::White, false),
        ));
        output.push_str(&format!(
            "{} {:.1}s\n",
            style_text("Runtime:", Color::Green, false),
            monitor.start_time.elapsed().as_secs_f64()
        ));
    }

    output.push('\n');

    // Charts using rasciichart
    let show_both = !args.download && !args.upload;

    if args.download || show_both {
        let dl_history = monitor.get_history_dl();
        if !dl_history.is_empty() {
            let chart = render_chart_rasciichart(
                &dl_history,
                args.height,
                chart_width,
                Color::Cyan,
                "▼ Download Speed",
            );
            output.push_str(&chart);
            output.push_str("\n\n");
        }
    }

    if (args.upload || show_both) && !args.download {
        let ul_history = monitor.get_history_ul();
        if !ul_history.is_empty() {
            let chart = render_chart_rasciichart(
                &ul_history,
                args.height,
                chart_width,
                Color::Yellow,
                "▲ Upload Speed",
            );
            output.push_str(&chart);
            output.push('\n');
        }
    }

    Ok(output)
}

fn monitor_bandwidth(args: Args) -> Result<()> {
    let interface = if let Some(iface) = args.iface.clone() {
        resolve_interface(&iface)?
    } else {
        select_best_interface()?
    };

    println!(
        "{} {}\n", 
        style_text("Monitoring interface:", Color::Green, false),
        style_text(&interface, Color::Cyan, true)
    );

    let monitor = Arc::new(Mutex::new(NetworkMonitor::new(interface, args.history)?));
    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();

    ctrlc::set_handler(move || {
        r.store(false, Ordering::SeqCst);
    })?;

    let mut stdout = stdout();
    execute!(stdout, EnterAlternateScreen, Hide)?;
    enable_raw_mode()?;

    let result = (|| -> Result<()> {
        let mut last_update = Instant::now();

        while running.load(Ordering::SeqCst) {
            // Check for key events (non-blocking)
            if event::poll(Duration::from_millis(50))? {
                if let Event::Key(key_event) = event::read()? {
                    match key_event.code {
                        KeyCode::Char('q') | KeyCode::Char('Q') | KeyCode::Esc => break,
                        KeyCode::Char('c') => {
                            use crossterm::event::KeyModifiers;
                            if key_event.modifiers.contains(KeyModifiers::CONTROL) {
                                break;
                            }
                        }
                        _ => {}
                    }
                }
            }

            // Update bandwidth stats with accurate timing
            if last_update.elapsed() >= INTERVAL {
                // Lock monitor untuk update - mencegah race condition
                let stats = {
                    let mut mon = monitor.lock().unwrap();
                    mon.update()?
                };

                // Get terminal size - ini bisa berubah karena resize
                let (term_width, term_height) = size()?;

                // Render UI dengan data terbaru
                let ui = {
                    let mon = monitor.lock().unwrap();
                    render_ui(&mon, &stats, &args, term_width)?
                };

                let mut lines: Vec<String> = ui.lines().map(str::to_owned).collect();

                // Resize output to fit terminal height
                lines.resize_with(term_height as usize, String::new);

                let full_output = lines.join("\n");

                // Write to screen
                execute!(
                    stdout,
                    MoveTo(0, 0),
                    Print(full_output)
                )?;
                stdout.flush()?;

                last_update = Instant::now();
            }
        }
        Ok(())
    })();

    // Cleanup
    disable_raw_mode()?;
    execute!(stdout, LeaveAlternateScreen, Show)?;

    if let Err(e) = result {
        eprintln!("{} {}", style_text("Error:", Color::Red, true), e);
    } else {
        println!("\n{}", style_text("Stopped cleanly.", Color::Green, true));
    }

    Ok(())
}

fn main() -> Result<()> {
    let args = Args::parse();

    if args.version {
        println!("{}", ColoredVersion::new());
        return Ok(());
    }

    if args.list {
        list_interfaces()?;
        return Ok(());
    }

    monitor_bandwidth(args)
}