use std::{
io,
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
use anyhow::Result;
use crossterm::{
event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
execute,
terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
};
use ratatui::{backend::CrosstermBackend, widgets::TableState, Terminal};
use crate::{
args::{Args, SortBy},
collector::{apply_filter, apply_sort, Collector},
types::Snapshot,
ui,
};
pub struct AppState {
pub snapshot: Snapshot,
pub args: Args,
pub cumulative: bool,
pub paused: bool,
pub filter_editing: bool,
pub filter_buf: String,
pub table_state: TableState,
pub row_count: usize,
pub show_help: bool,
pub start_time: u64, }
impl AppState {
pub fn scroll_down(&mut self) {
let max = self.row_count.saturating_sub(1);
let cur = self.table_state.selected().unwrap_or(0);
let next = (cur + 1).min(max);
self.table_state.select(Some(next));
}
pub fn scroll_up(&mut self) {
let cur = self.table_state.selected().unwrap_or(0);
let next = cur.saturating_sub(1);
self.table_state.select(Some(next));
}
pub fn jump_top(&mut self) {
self.table_state.select(Some(0));
}
pub fn jump_bottom(&mut self) {
let max = self.row_count.saturating_sub(1);
self.table_state.select(Some(max));
}
pub fn uptime_secs(&self) -> u64 {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
now.saturating_sub(self.start_time)
}
}
pub fn run(args: Args) -> Result<()> {
if args.json {
return run_json(args);
}
run_tui(args)
}
fn run_json(args: Args) -> Result<()> {
let mut collector = Collector::new();
std::thread::sleep(Duration::from_millis(args.interval));
let mut snap = collector.collect();
apply_filter(&mut snap.entries, &args.filter);
apply_sort(&mut snap.entries, &args.sort, args.cumulative);
if args.top > 0 {
snap.entries.truncate(args.top);
}
println!("{}", serde_json::to_string_pretty(&snap.entries)?);
Ok(())
}
fn run_tui(args: Args) -> Result<()> {
enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
let backend = CrosstermBackend::new(stdout);
let mut term = Terminal::new(backend)?;
term.hide_cursor()?;
let default_hook = std::panic::take_hook();
std::panic::set_hook(Box::new(move |info| {
let _ = disable_raw_mode();
let _ = execute!(io::stdout(), LeaveAlternateScreen);
default_hook(info);
}));
let mut collector = Collector::new();
let mut tick_count = 0usize;
let count_limit = args.count;
let interval = Duration::from_millis(args.interval);
std::thread::sleep(Duration::from_millis(200));
let first_snap = make_snapshot(&mut collector, &args);
let row_count = first_snap.entries.len();
let start_time = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let mut state = AppState {
snapshot: first_snap,
cumulative: args.cumulative,
paused: false,
filter_editing: false,
filter_buf: args.filter.clone().unwrap_or_default(),
table_state: TableState::default(),
row_count,
show_help: false,
start_time,
args,
};
state.table_state.select(Some(0));
let mut last_tick = Instant::now();
loop {
term.draw(|f| ui::draw(f, &state))?;
let timeout = interval
.checked_sub(last_tick.elapsed())
.unwrap_or(Duration::ZERO);
if event::poll(timeout)? {
if let Event::Key(key) = event::read()? {
if state.filter_editing {
handle_filter_key(&mut state, key);
} else if state.show_help {
state.show_help = false;
} else {
match handle_key(&mut state, key) {
KeyAction::Quit => break,
KeyAction::Continue => {}
}
}
}
}
if last_tick.elapsed() >= interval && !state.paused {
let snap = make_snapshot(&mut collector, &state.args);
state.row_count = snap.entries.len();
if let Some(sel) = state.table_state.selected() {
if sel >= state.row_count && state.row_count > 0 {
state.table_state.select(Some(state.row_count - 1));
}
}
state.snapshot = snap;
last_tick = Instant::now();
tick_count += 1;
if count_limit > 0 && tick_count >= count_limit {
break;
}
}
}
disable_raw_mode()?;
execute!(term.backend_mut(), LeaveAlternateScreen)?;
term.show_cursor()?;
Ok(())
}
fn make_snapshot(collector: &mut Collector, args: &Args) -> Snapshot {
let mut snap = collector.collect();
apply_filter(&mut snap.entries, &args.filter);
apply_sort(&mut snap.entries, &args.sort, args.cumulative);
if args.top > 0 {
snap.entries.truncate(args.top);
}
snap
}
enum KeyAction {
Quit,
Continue,
}
fn handle_key(state: &mut AppState, key: KeyEvent) -> KeyAction {
match key.code {
KeyCode::Char('q') | KeyCode::Char('Q') => return KeyAction::Quit,
KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => {
return KeyAction::Quit
}
KeyCode::Char('p') | KeyCode::Char('P') => state.paused = !state.paused,
KeyCode::Tab => {
state.cumulative = !state.cumulative;
state.args.cumulative = state.cumulative;
}
KeyCode::Char('s') | KeyCode::Char('S') => {
state.args.sort = next_sort(&state.args.sort);
}
KeyCode::Char('r') | KeyCode::Char('R') => {
for e in &mut state.snapshot.entries {
e.sent_total = 0;
e.recv_total = 0;
}
}
KeyCode::Char('f') | KeyCode::Char('F') => {
state.filter_editing = true;
state.filter_buf = state.args.filter.clone().unwrap_or_default();
}
KeyCode::Char('h') | KeyCode::Char('H') | KeyCode::Char('?') => {
state.show_help = !state.show_help;
}
KeyCode::Down | KeyCode::Char('j') => state.scroll_down(),
KeyCode::Up | KeyCode::Char('k') => state.scroll_up(),
KeyCode::Char('g') => state.jump_top(),
KeyCode::Char('G') => state.jump_bottom(),
KeyCode::Home => state.jump_top(),
KeyCode::End => state.jump_bottom(),
KeyCode::PageDown => {
for _ in 0..10 {
state.scroll_down();
}
}
KeyCode::PageUp => {
for _ in 0..10 {
state.scroll_up();
}
}
_ => {}
}
KeyAction::Continue
}
fn handle_filter_key(state: &mut AppState, key: KeyEvent) {
match key.code {
KeyCode::Enter => {
state.args.filter = if state.filter_buf.is_empty() {
None
} else {
Some(state.filter_buf.clone())
};
state.filter_editing = false;
}
KeyCode::Esc => {
state.filter_editing = false;
}
KeyCode::Backspace => {
state.filter_buf.pop();
}
KeyCode::Delete => {
state.filter_buf.clear();
}
KeyCode::Char(c) => {
state.filter_buf.push(c);
}
_ => {}
}
}
fn next_sort(current: &SortBy) -> SortBy {
match current {
SortBy::TotalRate => SortBy::Sent,
SortBy::Sent => SortBy::Recv,
SortBy::Recv => SortBy::SentTotal,
SortBy::SentTotal => SortBy::RecvTotal,
SortBy::RecvTotal => SortBy::Pid,
SortBy::Pid => SortBy::Name,
SortBy::Name => SortBy::TotalRate,
}
}