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},
    text::{Line, Span},
    widgets::{Block, Borders, Cell, Paragraph, Row, Table, TableState},
    Frame,
};
use sqlx::PgPool;

use crate::tui::app::{copy_to_clipboard, AppAction};

struct ProfiledQuery {
    index: usize,
    sql: String,
    duration_ms: u64,
    rows_affected: u64,
    table: String,
}

pub struct QueryProfilerTab {
    queries: Vec<ProfiledQuery>,
    list_state: TableState,
    detail_scroll: u16,
    explain_text: Option<String>,
    explain_cost: Option<f64>,
    explain_loading: bool,
    status: String,
}

impl Default for QueryProfilerTab {
    fn default() -> Self {
        let mut list_state = TableState::default();
        list_state.select(Some(0));
        Self {
            queries: Vec::new(),
            list_state,
            detail_scroll: 0,
            explain_text: None,
            explain_cost: None,
            explain_loading: false,
            status: "ready".into(),
        }
    }
}

impl QueryProfilerTab {
    pub async fn load(&mut self, _pool: &PgPool) {
        let raw = rok_orm::snapshot();
        self.queries = raw
            .into_iter()
            .rev()
            .enumerate()
            .map(|(i, e)| ProfiledQuery {
                index: i + 1,
                sql: e.sql,
                duration_ms: e.duration_ms,
                rows_affected: e.rows_affected,
                table: e.table,
            })
            .collect();
        let selected = self.list_state.selected().unwrap_or(0);
        if !self.queries.is_empty() && selected >= self.queries.len() {
            self.list_state.select(Some(self.queries.len().saturating_sub(1)));
        }
        self.status = format!("{} queries captured", self.queries.len());
        self.explain_loading = false;
    }

    pub fn render(&mut self, frame: &mut Frame, area: Rect) {
        let has_detail = self.queries.len() > 1
            || (!self.queries.is_empty() && self.queries[0].sql.len() > 80);

        let constraints = if has_detail {
            [Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)]
        } else {
            [Constraint::Min(0), Constraint::Length(0)]
        };

        let chunks = Layout::default()
            .direction(Direction::Vertical)
            .constraints(constraints)
            .split(area);

        self.render_query_list(frame, chunks[0]);
        self.render_detail(frame, chunks[1]);
    }

    fn render_query_list(&mut self, frame: &mut Frame, area: Rect) {
        let header = Row::new(vec![
            Cell::from("#").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("Duration").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("Rows").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("Table").style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("SQL").style(Style::default().add_modifier(Modifier::BOLD)),
        ])
        .height(1);

        let max_dur = self.queries.iter().map(|q| q.duration_ms).max().unwrap_or(1).max(1);

        let rows: Vec<Row> = self
            .queries
            .iter()
            .map(|q| {
                let dur_color = if q.duration_ms >= 100 {
                    Color::Red
                } else if q.duration_ms >= 10 {
                    Color::Yellow
                } else {
                    Color::Green
                };

                let bar_len = ((q.duration_ms as f64 / max_dur as f64) * 8.0) as usize;
                let bar = "â–“".repeat(bar_len);

                let sql_preview = if q.sql.len() > 60 {
                    format!("{}...", &q.sql[..57])
                } else {
                    q.sql.clone()
                };

                Row::new(vec![
                    Cell::from(q.index.to_string()),
                    Cell::from(format!("{}{}ms", bar, q.duration_ms))
                        .style(Style::default().fg(dur_color)),
                    Cell::from(q.rows_affected.to_string())
                        .style(Style::default().fg(Color::DarkGray)),
                    Cell::from(q.table.as_str()).style(Style::default().fg(Color::Cyan)),
                    Cell::from(sql_preview),
                ])
                .style(Style::default().bg(Color::Reset))
            })
            .collect();

        let widths = [
            Constraint::Length(5),
            Constraint::Length(10),
            Constraint::Length(6),
            Constraint::Length(14),
            Constraint::Min(30),
        ];

        let title = format!(" Query Profiler — {} ", self.status);
        let table = Table::new(rows, widths)
            .header(header)
            .block(
                Block::default()
                    .title(title)
                    .borders(Borders::ALL)
                    .border_style(Style::default().fg(Color::Cyan)),
            )
            .highlight_style(Style::default().bg(Color::DarkGray))
            .highlight_symbol("â–¶ ");

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

    fn render_detail(&mut self, frame: &mut Frame, area: Rect) {
        if area.width < 5 || area.height < 3 {
            return;
        }

        let selected = self.list_state.selected().and_then(|i| self.queries.get(i));
        let Some(query) = selected else {
            let para =
                Paragraph::new(" Select a query to see details ").style(Style::default().fg(Color::DarkGray)).block(
                    Block::default()
                        .title(" Detail ")
                        .borders(Borders::ALL),
                );
            frame.render_widget(para, area);
            return;
        };

        let chunks = Layout::default()
            .direction(Direction::Vertical)
            .constraints([Constraint::Min(3), Constraint::Length(3)])
            .split(area);

        let dur_color = if query.duration_ms >= 100 {
            Color::Red
        } else if query.duration_ms >= 10 {
            Color::Yellow
        } else {
            Color::Green
        };

        let mut lines = vec![
            Line::from(vec![
                Span::styled("SQL:      ", Style::default().fg(Color::Cyan)),
                Span::raw(&query.sql),
            ]),
            Line::from(vec![
                Span::styled("Duration: ", Style::default().fg(Color::Cyan)),
                Span::styled(
                    format!("{} ms", query.duration_ms),
                    Style::default().fg(dur_color),
                ),
                Span::raw("  "),
                Span::styled("Rows: ", Style::default().fg(Color::Cyan)),
                Span::raw(query.rows_affected.to_string()),
                Span::raw("  "),
                Span::styled("Table: ", Style::default().fg(Color::Cyan)),
                Span::raw(&query.table),
            ]),
        ];

        if let Some(ref explain) = self.explain_text {
            let cost_str = self
                .explain_cost
                .map(|c| format!(" (cost={:.2})", c))
                .unwrap_or_default();
            lines.push(Line::from(vec![
                Span::styled(
                    format!("Plan{cost_str}:"),
                    Style::default().fg(Color::Yellow),
                ),
            ]));
            for line in explain.lines() {
                lines.push(Line::from(Span::raw(line.to_string())));

                let suggestion = extract_suggestion(line);
                if let Some(s) = suggestion {
                    lines.push(Line::from(vec![
                        Span::styled("💡 ", Style::default().fg(Color::Green)),
                        Span::styled(s, Style::default().fg(Color::Green)),
                    ]));
                }
            }
        } else if self.explain_loading {
            lines.push(Line::from(vec![
                Span::styled(" Running EXPLAIN ANALYZE...", Style::default().fg(Color::Yellow)),
            ]));
        } else {
            lines.push(Line::from(vec![
                Span::styled(" Press ", Style::default().fg(Color::DarkGray)),
                Span::styled("p", Style::default().fg(Color::Cyan)),
                Span::styled(" to run EXPLAIN ANALYZE", Style::default().fg(Color::DarkGray)),
            ]));
        }

        let detail_w = chunks[0].width.max(5) as usize;
        let wrapped: Vec<Line> = lines
            .into_iter()
            .flat_map(|line| {
                let full = line_to_text(&line);
                let chars: Vec<char> = full.chars().collect();
                if chars.len() > detail_w {
                    chars
                        .chunks(detail_w)
                        .map(|chunk| Line::from(Span::raw(chunk.iter().collect::<String>())))
                        .collect::<Vec<_>>()
                } else {
                    vec![line]
                }
            })
            .collect();

        let para = Paragraph::new(wrapped)
            .scroll((self.detail_scroll, 0))
            .block(
                Block::default()
                    .title(" Detail ")
                    .borders(Borders::ALL),
            );
        frame.render_widget(para, chunks[0]);

        let help_text = Line::from(vec![
            Span::styled(" p ", Style::default().fg(Color::Cyan)),
            Span::raw("EXPLAIN  "),
            Span::styled(" c ", Style::default().fg(Color::Cyan)),
            Span::raw("Copy SQL  "),
            Span::styled(" r/F5 ", Style::default().fg(Color::Cyan)),
            Span::raw("Refresh  "),
            Span::styled(" ↑↓ ", Style::default().fg(Color::Cyan)),
            Span::raw("Navigate"),
        ]);
        let help = Paragraph::new(help_text)
            .style(Style::default().fg(Color::DarkGray))
            .block(
                Block::default()
                    .title(" Keys ")
                    .borders(Borders::ALL),
            );
        frame.render_widget(help, chunks[1]);
    }

    pub async fn explain_selected(&mut self, pool: &PgPool) {
        let Some(idx) = self.list_state.selected() else { return };
        let Some(query) = self.queries.get(idx) else { return };
        self.explain_loading = true;

        let explain_sql = format!("EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) {}", query.sql);
        let result = sqlx::query(&explain_sql)
            .fetch_all(pool)
            .await;

        match result {
            Ok(rows) => {
                use sqlx::Row as _;
                let lines: Vec<String> = rows
                    .iter()
                    .map(|r| r.try_get::<String, _>(0).unwrap_or_default())
                    .collect();
                let plan_text = lines.join("\n");
                let cost: f64 = lines
                    .first()
                    .and_then(|l| {
                        l.split("cost=")
                            .nth(1)
                            .and_then(|s| s.split("..").nth(1))
                            .and_then(|s| s.split_whitespace().next())
                            .and_then(|s| s.parse().ok())
                    })
                    .unwrap_or(0.0);
                self.explain_text = Some(plan_text);
                self.explain_cost = Some(cost);
                self.status = "EXPLAIN ANALYZE complete".into();
            }
            Err(e) => {
                self.explain_text = Some(format!("Error: {e}"));
                self.explain_cost = None;
                self.status = "EXPLAIN failed".into();
            }
        }
        self.explain_loading = false;
    }

    pub fn handle_key(&mut self, key: KeyEvent) -> AppAction {
        match key.code {
            KeyCode::Up => {
                let i = self.list_state.selected().unwrap_or(0).saturating_sub(1);
                self.list_state.select(Some(i));
                AppAction::None
            }
            KeyCode::Down => {
                let max = self.queries.len().saturating_sub(1);
                let i = (self.list_state.selected().unwrap_or(0) + 1).min(max);
                self.list_state.select(Some(i));
                AppAction::None
            }
            KeyCode::Char('c') | KeyCode::Char('C') => {
                if let Some(idx) = self.list_state.selected() {
                    if let Some(q) = self.queries.get(idx) {
                        copy_to_clipboard(&q.sql);
                        return AppAction::Reload;
                    }
                }
                AppAction::None
            }
            KeyCode::Char('p') | KeyCode::Char('P') => {
                AppAction::ProfileQuery
            }
            KeyCode::F(5) | KeyCode::Char('r') | KeyCode::Char('R') => AppAction::Reload,
            _ => AppAction::None,
        }
    }
}

fn line_to_text(line: &Line) -> String {
    line.spans.iter().map(|s| &*s.content).collect()
}

fn extract_suggestion(line: &str) -> Option<String> {
    let lower = line.to_lowercase();
    if lower.contains("seq scan") {
        if let Some(col_part) = line.split("on ").nth(1) {
            let table_hint = col_part
                .split_whitespace()
                .next()
                .unwrap_or("table")
                .trim_matches(|c: char| c.is_ascii_punctuation());
            return Some(format!(
                "Consider adding an index on filtered columns of `{table_hint}`"
            ));
        }
        return Some("Consider adding an index on filtered columns".into());
    }
    None
}