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;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DbPanel {
    #[default]
    TableList,
    RowGrid,
    SqlPane,
}

pub struct DbTab {
    pub tables: Vec<String>,
    pub table_state: TableState,
    pub columns: Vec<String>,
    pub rows: Vec<Vec<String>>,
    pub row_state: TableState,
    pub page: usize,
    pub total: usize,
    pub per_page: usize,
    pub panel: DbPanel,
    pub sql_buf: String,
    pub status: String,
}

impl Default for DbTab {
    fn default() -> Self {
        Self {
            tables: Vec::new(),
            table_state: TableState::default(),
            columns: Vec::new(),
            rows: Vec::new(),
            row_state: TableState::default(),
            page: 0,
            total: 0,
            per_page: 50,
            panel: DbPanel::TableList,
            sql_buf: String::new(),
            status: "Loading...".into(),
        }
    }
}

impl DbTab {
    pub async fn load(&mut self, pool: &PgPool) {
        match list_tables(pool).await {
            Ok(tables) => {
                self.tables = tables;
                if self.table_state.selected().is_none() && !self.tables.is_empty() {
                    self.table_state.select(Some(0));
                    self.load_rows(pool).await;
                }
                self.status = format!("{} tables", self.tables.len());
            }
            Err(e) => {
                self.status = format!("Error: {e}");
            }
        }
    }

    pub async fn load_rows(&mut self, pool: &PgPool) {
        let Some(idx) = self.table_state.selected() else {
            return;
        };
        let Some(table) = self.tables.get(idx) else {
            return;
        };
        let table = table.clone();
        let offset = self.page * self.per_page;

        let sql = if self.sql_buf.trim().is_empty() {
            format!(
                "SELECT * FROM \"{table}\" LIMIT {} OFFSET {}",
                self.per_page, offset
            )
        } else {
            self.sql_buf.clone()
        };

        match fetch_rows(pool, &sql).await {
            Ok((cols, rows)) => {
                self.columns = cols;
                self.rows = rows;
                self.row_state
                    .select(if self.rows.is_empty() { None } else { Some(0) });
                self.status = format!("Table: {table}  page {}", self.page + 1);
            }
            Err(e) => {
                self.status = format!("Query error: {e}");
                self.columns.clear();
                self.rows.clear();
            }
        }
    }

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

        let top = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Length(22), Constraint::Min(10)])
            .split(chunks[0]);

        self.render_table_list(frame, top[0]);
        self.render_row_grid(frame, top[1]);
        self.render_sql_pane(frame, chunks[1]);
    }

    fn render_table_list(&mut self, frame: &mut Frame, area: Rect) {
        let focused = self.panel == DbPanel::TableList;
        let border_style = if focused {
            Style::default().fg(Color::Cyan)
        } else {
            Style::default()
        };

        let rows: Vec<Row> = self
            .tables
            .iter()
            .map(|t| Row::new(vec![Cell::from(t.as_str())]))
            .collect();

        let widths = [Constraint::Percentage(100)];
        let table = Table::new(rows, widths)
            .block(
                Block::default()
                    .title(" Tables ")
                    .borders(Borders::ALL)
                    .border_style(border_style),
            )
            .highlight_style(
                Style::default()
                    .fg(Color::Black)
                    .bg(Color::Cyan)
                    .add_modifier(Modifier::BOLD),
            )
            .highlight_symbol("â–¶ ");

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

    fn render_row_grid(&mut self, frame: &mut Frame, area: Rect) {
        let focused = self.panel == DbPanel::RowGrid;
        let border_style = if focused {
            Style::default().fg(Color::Cyan)
        } else {
            Style::default()
        };

        let title = format!(" {} ", self.status);

        if self.columns.is_empty() {
            let p = Paragraph::new("No rows").block(
                Block::default()
                    .title(title)
                    .borders(Borders::ALL)
                    .border_style(border_style),
            );
            frame.render_widget(p, area);
            return;
        }

        let header: Vec<Cell> = self
            .columns
            .iter()
            .map(|c| Cell::from(c.as_str()).style(Style::default().add_modifier(Modifier::BOLD)))
            .collect();
        let header_row = Row::new(header).height(1).bottom_margin(0);

        let rows: Vec<Row> = self
            .rows
            .iter()
            .map(|row| {
                let cells: Vec<Cell> = row.iter().map(|v| Cell::from(v.as_str())).collect();
                Row::new(cells)
            })
            .collect();

        let col_count = self.columns.len().max(1);
        let widths: Vec<Constraint> = (0..col_count).map(|_| Constraint::Min(12)).collect();

        let table = Table::new(rows, widths)
            .header(header_row)
            .block(
                Block::default()
                    .title(title)
                    .borders(Borders::ALL)
                    .border_style(border_style),
            )
            .highlight_style(Style::default().bg(Color::DarkGray));

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

    fn render_sql_pane(&mut self, frame: &mut Frame, area: Rect) {
        let focused = self.panel == DbPanel::SqlPane;
        let border_style = if focused {
            Style::default().fg(Color::Yellow)
        } else {
            Style::default()
        };

        let display = if self.sql_buf.is_empty() {
            "[Q] Enter SQL — Tab to switch panels — ↑↓ Navigate — F5 Refresh".to_string()
        } else {
            self.sql_buf.clone()
        };

        let p = Paragraph::new(display).block(
            Block::default()
                .title(" SQL Query ")
                .borders(Borders::ALL)
                .border_style(border_style),
        );

        frame.render_widget(p, area);
    }

    pub fn selected_row_text(&self) -> String {
        if self.columns.is_empty() || self.rows.is_empty() {
            return String::new();
        }
        if let Some(selected) = self.row_state.selected() {
            if selected < self.rows.len() {
                let row = &self.rows[selected];
                let pairs: Vec<String> = self
                    .columns
                    .iter()
                    .enumerate()
                    .map(|(i, col)| {
                        format!(
                            "{}: {}",
                            col,
                            row.get(i).map(|s| s.as_str()).unwrap_or("NULL")
                        )
                    })
                    .collect();
                return pairs.join("\n");
            }
        }
        let rows_json: Vec<serde_json::Value> = self
            .rows
            .iter()
            .map(|row| {
                let mut map = serde_json::Map::new();
                for (i, col) in self.columns.iter().enumerate() {
                    map.insert(
                        col.clone(),
                        serde_json::Value::String(row.get(i).cloned().unwrap_or_default()),
                    );
                }
                serde_json::Value::Object(map)
            })
            .collect();
        serde_json::to_string_pretty(&rows_json).unwrap_or_default()
    }

    pub fn handle_key(&mut self, key: KeyEvent) -> AppAction {
        match key.code {
            KeyCode::Tab => {
                self.panel = match self.panel {
                    DbPanel::TableList => DbPanel::RowGrid,
                    DbPanel::RowGrid => DbPanel::SqlPane,
                    DbPanel::SqlPane => DbPanel::TableList,
                };
                AppAction::None
            }
            KeyCode::F(5) => AppAction::Reload,
            KeyCode::Up => match self.panel {
                DbPanel::TableList => {
                    let i = self.table_state.selected().unwrap_or(0).saturating_sub(1);
                    self.table_state.select(Some(i));
                    AppAction::LoadRows
                }
                DbPanel::RowGrid => {
                    let i = self.row_state.selected().unwrap_or(0).saturating_sub(1);
                    self.row_state.select(Some(i));
                    AppAction::None
                }
                _ => AppAction::None,
            },
            KeyCode::Down => match self.panel {
                DbPanel::TableList => {
                    let max = self.tables.len().saturating_sub(1);
                    let i = (self.table_state.selected().unwrap_or(0) + 1).min(max);
                    self.table_state.select(Some(i));
                    AppAction::LoadRows
                }
                DbPanel::RowGrid => {
                    let max = self.rows.len().saturating_sub(1);
                    let i = (self.row_state.selected().unwrap_or(0) + 1).min(max);
                    self.row_state.select(Some(i));
                    AppAction::None
                }
                _ => AppAction::None,
            },
            KeyCode::PageUp => {
                self.page = self.page.saturating_sub(1);
                AppAction::LoadRows
            }
            KeyCode::PageDown => {
                self.page += 1;
                AppAction::LoadRows
            }
            KeyCode::Char('q') | KeyCode::Char('Q') if self.panel != DbPanel::SqlPane => {
                self.panel = DbPanel::SqlPane;
                self.sql_buf.clear();
                AppAction::None
            }
            KeyCode::Char(c) if self.panel == DbPanel::SqlPane => {
                self.sql_buf.push(c);
                AppAction::None
            }
            KeyCode::Backspace if self.panel == DbPanel::SqlPane => {
                self.sql_buf.pop();
                AppAction::None
            }
            KeyCode::Enter if self.panel == DbPanel::SqlPane => AppAction::LoadRows,
            KeyCode::Esc => {
                if self.panel == DbPanel::SqlPane {
                    self.sql_buf.clear();
                    self.panel = DbPanel::TableList;
                }
                AppAction::None
            }
            _ => AppAction::None,
        }
    }
}

async fn list_tables(pool: &PgPool) -> anyhow::Result<Vec<String>> {
    let rows = sqlx::query_scalar::<_, String>(
        "SELECT table_name FROM information_schema.tables \
         WHERE table_schema = 'public' AND table_type = 'BASE TABLE' \
         ORDER BY table_name",
    )
    .fetch_all(pool)
    .await?;
    Ok(rows)
}

async fn fetch_rows(
    pool: &PgPool,
    sql: &str,
) -> anyhow::Result<(Vec<String>, Vec<Vec<String>>)> {
    use sqlx::{Column as _, Row as _};

    let rows = sqlx::query(sql).fetch_all(pool).await?;

    if rows.is_empty() {
        return Ok((Vec::new(), Vec::new()));
    }

    let columns: Vec<String> = rows[0]
        .columns()
        .iter()
        .map(|c| c.name().to_string())
        .collect();

    let data: Vec<Vec<String>> = rows
        .iter()
        .map(|row| {
            columns
                .iter()
                .enumerate()
                .map(|(i, _)| {
                    row.try_get::<String, _>(i)
                        .or_else(|_| row.try_get::<i64, _>(i).map(|n| n.to_string()))
                        .or_else(|_| row.try_get::<bool, _>(i).map(|b| b.to_string()))
                        .or_else(|_| row.try_get::<f64, _>(i).map(|f| f.to_string()))
                        .unwrap_or_else(|_| "NULL".to_string())
                })
                .collect()
        })
        .collect();

    Ok((columns, data))
}