rok-cli 0.6.1

Developer CLI for rok-based Axum applications
use crossterm::event::{KeyCode, KeyEvent};
use ratatui::{
    layout::{Constraint, Direction, Layout, Rect},
    style::{Color, Modifier, Style},
    widgets::{Block, Borders, Cell, Paragraph, Row, Table, TableState},
    Frame,
};
use sqlx::PgPool;

use crate::tui::app::AppAction;

pub struct SlowQuery {
    pub sql: String,
    pub duration_ms: f64,
    pub state: String,
    pub pid: i32,
}

pub struct PoolStats {
    pub total: i64,
    pub active: i64,
    pub idle: i64,
    pub waiting: i64,
}

pub struct PerfTab {
    pub slow_queries: Vec<SlowQuery>,
    pub pool_stats: Option<PoolStats>,
    pub slow_state: TableState,
    pub status: String,
}

impl Default for PerfTab {
    fn default() -> Self {
        let mut slow_state = TableState::default();
        slow_state.select(Some(0));
        Self {
            slow_queries: Vec::new(),
            pool_stats: None,
            slow_state,
            status: "Loading...".into(),
        }
    }
}

impl PerfTab {
    pub async fn load(&mut self, pool: &PgPool) {
        let (slow, stats) = fetch_perf_data(pool).await;
        self.slow_queries = slow;
        self.pool_stats = Some(stats);
        self.status = format!("{} slow queries", self.slow_queries.len());
    }

    pub fn render(&mut self, frame: &mut Frame, area: Rect) {
        let chunks = Layout::default()
            .direction(Direction::Vertical)
            .constraints([Constraint::Length(5), Constraint::Min(5)])
            .split(area);

        self.render_pool_stats(frame, chunks[0]);
        self.render_slow_queries(frame, chunks[1]);
    }

    fn render_pool_stats(&self, frame: &mut Frame, area: Rect) {
        let text = match &self.pool_stats {
            Some(s) => format!(
                " Total: {}  Active: {}  Idle: {}  Waiting: {}",
                s.total, s.active, s.idle, s.waiting
            ),
            None => " Loading pool stats...".to_string(),
        };

        let para = Paragraph::new(text)
            .style(Style::default().fg(Color::Cyan))
            .block(
                Block::default()
                    .title(" DB Pool (pg_stat_activity) ")
                    .borders(Borders::ALL),
            );
        frame.render_widget(para, area);
    }

    fn render_slow_queries(&mut self, frame: &mut Frame, area: Rect) {
        let header = Row::new(vec![
            Cell::from("PID").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("Duration (ms)").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("State").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("Query").style(Style::default().add_modifier(Modifier::BOLD)),
        ])
        .height(1);

        let rows: Vec<Row> = self
            .slow_queries
            .iter()
            .map(|q| {
                let dur_color = if q.duration_ms > 1000.0 {
                    Color::Red
                } else if q.duration_ms > 300.0 {
                    Color::Yellow
                } else {
                    Color::White
                };
                Row::new(vec![
                    Cell::from(q.pid.to_string()),
                    Cell::from(format!("{:.1}", q.duration_ms))
                        .style(Style::default().fg(dur_color)),
                    Cell::from(q.state.as_str()).style(Style::default().fg(Color::DarkGray)),
                    Cell::from(q.sql.as_str()),
                ])
            })
            .collect();

        let widths = [
            Constraint::Length(7),
            Constraint::Length(14),
            Constraint::Length(12),
            Constraint::Min(30),
        ];

        let title = format!(" Slow Queries (top 10) — {} ", self.status);
        let table = Table::new(rows, widths)
            .header(header)
            .block(
                Block::default()
                    .title(title)
                    .borders(Borders::ALL)
                    .border_style(Style::default().fg(Color::Yellow)),
            )
            .highlight_style(Style::default().bg(Color::DarkGray))
            .highlight_symbol("");

        frame.render_stateful_widget(table, area, &mut self.slow_state);
    }

    pub fn handle_key(&mut self, key: KeyEvent) -> AppAction {
        match key.code {
            KeyCode::Up => {
                let i = self.slow_state.selected().unwrap_or(0).saturating_sub(1);
                self.slow_state.select(Some(i));
                AppAction::None
            }
            KeyCode::Down => {
                let max = self.slow_queries.len().saturating_sub(1);
                let i = (self.slow_state.selected().unwrap_or(0) + 1).min(max);
                self.slow_state.select(Some(i));
                AppAction::None
            }
            KeyCode::F(5) | KeyCode::Char('r') | KeyCode::Char('R') => AppAction::Reload,
            _ => AppAction::None,
        }
    }
}

async fn fetch_perf_data(pool: &PgPool) -> (Vec<SlowQuery>, PoolStats) {
    use sqlx::Row as _;

    let slow = sqlx::query(
        "SELECT pid, \
                EXTRACT(EPOCH FROM (now() - query_start)) * 1000 AS duration_ms, \
                state, \
                LEFT(query, 200) AS query \
         FROM pg_stat_activity \
         WHERE state != 'idle' \
           AND query_start IS NOT NULL \
           AND query NOT LIKE '%pg_stat_activity%' \
         ORDER BY duration_ms DESC NULLS LAST \
         LIMIT 10",
    )
    .fetch_all(pool)
    .await
    .unwrap_or_default()
    .into_iter()
    .map(|r| SlowQuery {
        pid: r.try_get::<i32, _>("pid").unwrap_or(0),
        duration_ms: r.try_get::<f64, _>("duration_ms").unwrap_or(0.0),
        state: r.try_get::<String, _>("state").unwrap_or_default(),
        sql: r.try_get::<String, _>("query").unwrap_or_default(),
    })
    .collect();

    let stats_row = sqlx::query(
        "SELECT COUNT(*) AS total, \
                COUNT(*) FILTER (WHERE state = 'active')  AS active, \
                COUNT(*) FILTER (WHERE state = 'idle')    AS idle, \
                COUNT(*) FILTER (WHERE wait_event IS NOT NULL AND state = 'active') AS waiting \
         FROM pg_stat_activity",
    )
    .fetch_one(pool)
    .await;

    let stats = match stats_row {
        Ok(r) => PoolStats {
            total: r.try_get("total").unwrap_or(0),
            active: r.try_get("active").unwrap_or(0),
            idle: r.try_get("idle").unwrap_or(0),
            waiting: r.try_get("waiting").unwrap_or(0),
        },
        Err(_) => PoolStats {
            total: 0,
            active: 0,
            idle: 0,
            waiting: 0,
        },
    };

    (slow, stats)
}