monothink 0.4.0

/*
 * This file is part of monothink.
 *
 * monothink is free software: you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free Software Foundation,
 * either version 3 of the License, or (at your option) any later version.
 *
 * monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with Wavething.
 * If not, see <https://www.gnu.org/licenses/>.
 */
use ratatui::{
    prelude::*,
    widgets::{Paragraph, Widget},
};

use crate::session::TokenUsage;

#[derive(Debug, Clone, Copy)]
pub enum InputMode {
    Normal,
    Insert,
    View,
    Command,
}

#[derive(Debug, Clone, Copy, Default)]
pub enum ActivityState {
    #[default]
    Idle,
    Active(usize),
}

const BARCOLOR: Color = Color::Rgb(0, 122, 204);

impl InputMode {
    fn label(self) -> &'static str {
        match self {
            Self::Normal | Self::Command => "NORMAL",
            Self::Insert => "INSERT",
            Self::View => "VIEW",
        }
    }

    fn colors(self) -> (Color, Color) {
        match self {
            Self::Normal | Self::Command => (Color::White, Color::Blue),
            Self::Insert => (Color::White, Color::Yellow),
            Self::View => (Color::White, Color::Magenta),
        }
    }
}

#[derive(Debug, Clone)]
pub struct StatusBarConfig {
    // @todo background color
    // @todo text color
    // @todo mode colors
    pub idle_symbol: String,
    pub activity_symbols: Vec<String>,
    pub input_tokens_symbol: String,
    pub output_tokens_symbol: String,
    pub leader_symbol: String,
    pub separator_symbol: String,
    pub reverse_separator_symbol: String,
}

#[derive(Debug, Clone)]
pub struct StatusBar {
    pub acitivity_state: ActivityState,
    pub config: StatusBarConfig,
    pub input_mode: InputMode,
    pub model: String,
    pub context_used: u64,
    pub context_total: u64,
    pub current_line: u64,
    pub total_lines: u64,
    pub tokens: TokenUsage,
    pub status_message: String,
}

impl Widget for StatusBar {
    fn render(self, area: Rect, buf: &mut Buffer) {
        let (mode_fg, mode_bg) = self.input_mode.colors();

        let ctx_percent = (self.context_used * 100)
            .checked_div(self.context_total)
            .unwrap_or_default();

        let left_spans = vec![
            // Input mode
            Span::styled(
                format!(" {} ", self.input_mode.label()),
                Style::default()
                    .fg(mode_fg)
                    .bg(mode_bg)
                    .add_modifier(Modifier::BOLD),
            ),
            Span::styled(
                &self.config.leader_symbol,
                Style::default().fg(mode_bg).bg(BARCOLOR),
            ),
            // Model
            Span::styled(
                format!(" {} ", self.model),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            Span::styled(
                &self.config.separator_symbol,
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            // Context
            Span::styled(
                // @todo adapt to kilo or millions depending on the context size
                format!(
                    "{:>3}% ({:>3}k of {:>3}k)",
                    ctx_percent,
                    self.context_used / 1000,
                    self.context_total / 1000
                ),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            Span::styled(
                &self.config.separator_symbol,
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            // Tokens
            Span::styled(
                format!(
                    " {}{} {}{} ",
                    self.config.input_tokens_symbol,
                    self.render_tokens(self.tokens.input),
                    self.config.output_tokens_symbol,
                    self.render_tokens(self.tokens.output)
                ),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            Span::styled(
                &self.config.separator_symbol,
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            // Assistant mode
            Span::styled(
                format!(" {} ", self.status_message),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
        ];

        let right_spans = vec![
            // acitvity
            Span::styled(
                format!(" {} ", self.acitivity_symbol()),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            Span::styled(
                &self.config.reverse_separator_symbol,
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
            // lines
            Span::styled(
                format!(" {} of {} ", self.current_line, self.total_lines),
                Style::default().fg(Color::White).bg(BARCOLOR),
            ),
        ];

        let [left_area, right_area] =
            Layout::horizontal([Constraint::Fill(1), Constraint::Length(20)]).areas(area);
        Paragraph::new(Line::from(left_spans))
            .style(Style::default().bg(BARCOLOR))
            .render(left_area, buf);
        Paragraph::new(Line::from(right_spans))
            .alignment(HorizontalAlignment::Right)
            .style(Style::default().bg(BARCOLOR))
            .render(right_area, buf);
    }
}

impl StatusBar {
    // if not acitivity symbols are set, we always show the idle symbols
    // This allows the user to disable the activity symbol
    pub fn indicate_activity(&mut self) {
        if self.config.activity_symbols.is_empty() {
            self.acitivity_state = ActivityState::Idle;
        } else {
            match self.acitivity_state {
                ActivityState::Idle => self.acitivity_state = ActivityState::Active(0),
                ActivityState::Active(index) => {
                    if index + 1 >= self.config.activity_symbols.len() {
                        self.acitivity_state = ActivityState::Active(0)
                    } else {
                        self.acitivity_state = ActivityState::Active(index + 1);
                    }
                }
            }
        }
    }

    pub fn indicate_idle(&mut self) {
        self.acitivity_state = ActivityState::Idle;
    }

    fn acitivity_symbol(&self) -> &str {
        match self.acitivity_state {
            ActivityState::Idle => &self.config.idle_symbol,
            ActivityState::Active(index) => self
                .config
                .activity_symbols
                .get(index)
                // fall back to idle to avoid crashing the application
                // hitting this case is an error in state update
                .unwrap_or(&self.config.idle_symbol),
        }
    }

    fn render_tokens(&self, num_tokens: Option<u64>) -> String {
        if let Some(num) = num_tokens {
            if num > 1_000_000_000 {
                format!("{:>3}B", num / 1_000_000_000)
            } else if num > 1_000_000 {
                format!("{:>3}M", num / 1_000_000)
            } else if num > 1_000 {
                format!("{:>3}k", num / 1000)
            } else {
                format!("{:>3}", num)
            }
        } else {
            "-".to_string()
        }
    }
}