rok-cli 0.6.1

Developer CLI for rok-based Axum applications
use std::collections::VecDeque;
use std::sync::{Mutex, OnceLock};

use chrono::{DateTime, Utc};
use ratatui::{
    layout::Rect,
    style::{Color, Modifier, Style},
    text::{Line, Span},
    widgets::{Block, Borders, Paragraph},
    Frame,
};
use tracing_subscriber::layer::Layer;

/// A single error/warn/slow-query entry in the waterfall.
#[derive(Debug, Clone)]
pub struct ErrorEntry {
    pub timestamp: DateTime<Utc>,
    pub level: String,
    pub message: String,
    pub source: String,
    pub count: u32,
}

/// Ring buffer for error entries.
pub struct ErrorWaterfall {
    entries: VecDeque<ErrorEntry>,
    capacity: usize,
}

impl ErrorWaterfall {
    pub fn new(capacity: usize) -> Self {
        Self {
            entries: VecDeque::with_capacity(capacity),
            capacity,
        }
    }

    pub fn push(&mut self, entry: ErrorEntry) {
        if let Some(last) = self.entries.back_mut() {
            let within_5s =
                (entry.timestamp - last.timestamp).num_seconds().abs() < 5;
            if within_5s
                && last.message == entry.message
                && last.source == entry.source
                && last.level == entry.level
            {
                last.count += 1;
                last.timestamp = entry.timestamp;
                return;
            }
        }

        if self.entries.len() >= self.capacity {
            self.entries.pop_front();
        }
        self.entries.push_back(entry);
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn latest(&self, n: usize) -> Vec<&ErrorEntry> {
        self.entries.iter().rev().take(n).collect()
    }
}

fn waterfall() -> &'static Mutex<ErrorWaterfall> {
    static WF: OnceLock<Mutex<ErrorWaterfall>> = OnceLock::new();
    WF.get_or_init(|| Mutex::new(ErrorWaterfall::new(100)))
}

/// Return the current number of entries in the waterfall.
pub fn waterfall_count() -> usize {
    waterfall().lock().map(|wf| wf.len()).unwrap_or(0)
}

/// Render the error waterfall pane.
pub fn render_waterfall(frame: &mut Frame, area: Rect) {
    let entries = waterfall().lock().unwrap();
    let count = entries.len();
    if count == 0 {
        return;
    }

    let lines: Vec<Line> = entries
        .latest(5)
        .iter()
        .map(|e| {
            let ts = e.timestamp.format("%H:%M:%S").to_string();
            let level_color = match e.level.as_str() {
                "ERROR" => Color::Red,
                "WARN" => Color::Yellow,
                _ => Color::Magenta,
            };

            let count_suffix = if e.count > 1 {
                format!("{})", e.count)
            } else {
                String::new()
            };

            let msg = if e.message.len() > 80 {
                format!("{}...", &e.message[..77])
            } else {
                e.message.clone()
            };

            Line::from(vec![
                Span::styled(
                    format!(" [{ts}]"),
                    Style::default().fg(Color::DarkGray),
                ),
                Span::raw(" "),
                Span::styled(
                    format!(" {:5}", e.level),
                    Style::default()
                        .fg(level_color)
                        .add_modifier(Modifier::BOLD),
                ),
                Span::raw(" "),
                Span::styled(msg, Style::default().fg(Color::White)),
                Span::styled(count_suffix, Style::default().fg(Color::Yellow)),
                Span::styled(
                    format!(" ({})", e.source),
                    Style::default().fg(Color::DarkGray),
                ),
            ])
        })
        .collect();

    let block = Block::default()
        .borders(Borders::TOP)
        .style(Style::default().bg(Color::Black));
    let inner = block.inner(area);
    for (i, line) in lines.iter().enumerate() {
        if i as u16 >= inner.height {
            break;
        }
        let line_area = Rect {
            x: inner.x,
            y: inner.y + i as u16,
            width: inner.width,
            height: 1,
        };
        frame.render_widget(Paragraph::new(line.clone()), line_area);
    }
}

/// The tracing subscriber layer that captures events into the waterfall.
pub struct ErrorCollector;

impl<S: tracing::Subscriber> Layer<S> for ErrorCollector {
    fn on_event(
        &self,
        event: &tracing::Event<'_>,
        _ctx: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let meta = event.metadata();
        if *meta.level() < tracing::Level::WARN {
            return;
        }

        let mut message = String::new();
        let mut visitor = MessageVisitor(&mut message);
        event.record(&mut visitor);

        if message.is_empty() {
            return;
        }

        let source = meta.module_path().unwrap_or("unknown");

        let entry = ErrorEntry {
            timestamp: Utc::now(),
            level: meta.level().to_string(),
            message,
            source: source.to_string(),
            count: 1,
        };

        if let Ok(mut wf) = waterfall().lock() {
            wf.push(entry);
        }
    }
}

struct MessageVisitor<'a>(&'a mut String);

impl<'a> tracing::field::Visit for MessageVisitor<'a> {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        if field.name() == "message" {
            let s = format!("{value:?}");
            self.0.push_str(s.trim_matches('"'));
        }
    }

    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        if field.name() == "message" {
            self.0.push_str(value);
        }
    }
}

/// Initialize the error collector as a global tracing subscriber.
pub fn init_error_collector() {
    use tracing_subscriber::layer::SubscriberExt;
    use tracing_subscriber::Registry;

    let subscriber = Registry::default().with(ErrorCollector);
    let _ = tracing::subscriber::set_global_default(subscriber);
}