monothink 0.4.2

/*
 * 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 crate::{
    config::{Symbols, ViewMode},
    core::conversation::{Block, ConversationItem, ConversationItemKind},
    session::{SessionMessage, TokenUsage},
    ui::{
        conversation::{ConversationView, ConversationViewModel, HighlightConfig},
        statusline::{InputMode, StatusBar},
    },
};

use anyhow::Result;
use copypasta::{ClipboardContext, ClipboardProvider};
use crossterm::event::{Event, KeyCode};
use crossterm::event::{EventStream, KeyModifiers};
use futures::StreamExt;
use ratatui::{
    Frame, Terminal,
    layout::{Constraint, Layout},
    prelude::CrosstermBackend,
    style::Style,
};
use ratatui_textarea::{CursorMove, TextArea, WrapMode};
use std::io::Stdout;
use tokio::{
    sync::mpsc::{Receiver, Sender},
    time::{Duration, interval},
};

pub mod conversation;
pub mod statusline;

const REDRAW_TICK: u64 = 20; // in ms
const ACTIVITY_REFRESH_INTERVAL: u64 = 10; // multiple of REDRAW_TICK

#[derive(Clone, Debug)]
pub enum ModalAction {
    CommandMode(char),
    Goto(u32),
    InsertMode,
    Invalid(String),
    MoveDown,
    MoveBack,
    MoveForward,
    MoveUp,
    MoveWordBack,
    MoveWordEnd,
    MoveWordForward,
    Paste,
    Redo,
    Undo,
    ViewMode,
    YankSelection,
    YankSnippet(u32),
}
#[derive(Clone, Debug)]
pub enum UiMessage {
    ActiveModel(Option<String>),
    ContextState { fill_state: u64, total: u64 },
    ConversationItem(ConversationItem),
    ConversationViewMode(ViewMode),
    GenerationFinished(TokenUsage),
    IsProcessing(bool),
    //GenerationStarted,
    //Error(String),
    Quit,
}

pub struct Ui {
    conversation: ConversationViewModel, // partially moves to Session
    input: TextArea<'static>,
    processing: bool,
    status: StatusBar,
    session_tx: Sender<SessionMessage>,
    rx: Receiver<UiMessage>,
    dirty: bool,
    quit: bool,
    clipboard: ClipboardContext,
    action: String,
    stashed_input: Vec<String>,
}

impl Ui {
    pub fn new(
        session_tx: Sender<SessionMessage>,
        rx: Receiver<UiMessage>,
        symbols: Symbols,
    ) -> Self {
        let mut input = TextArea::default();
        input.set_cursor_line_style(Style::default());
        input.set_wrap_mode(WrapMode::Word);

        Ui {
            conversation: ConversationViewModel::new(
                HighlightConfig::new(),
                conversation::SymbolConfig {
                    assistant_message: symbols.assistant_message,
                    assistant_reasoning: symbols.assistant_reasoning,
                    tool_message: symbols.tool_message,
                    message_state_final: symbols.message_state_final,
                    message_state_intermediate: symbols.message_state_intermediate,
                    message_notification_debug: symbols.message_notification_debug,
                    message_notification_error: symbols.message_notification_error,
                    message_notification_info: symbols.message_notification_info,
                    message_notification_warning: symbols.message_notification_warning,
                    status_failed: symbols.status_failed,
                    status_success: symbols.status_success,
                    system_message: symbols.system_message,
                    user_message: symbols.user_message,
                },
            ),
            input,
            processing: false,
            status: StatusBar {
                acitivity_state: statusline::ActivityState::Idle,
                config: statusline::StatusBarConfig {
                    idle_symbol: symbols.idle_indicator,
                    activity_symbols: symbols.active_indicators,
                    input_tokens_symbol: symbols.input_tokens,
                    output_tokens_symbol: symbols.output_tokens,
                    leader_symbol: symbols.statusbar_leader,
                    separator_symbol: symbols.statusbar_separator,
                    reverse_separator_symbol: symbols.statusbar_reverse_separator,
                },
                input_mode: InputMode::Normal,
                model: "".to_string(),
                context_used: 0,
                context_total: 0,
                current_line: 0,
                total_lines: 0,
                tokens: TokenUsage::default(),
                status_message: String::new(),
            },
            session_tx,
            rx,
            dirty: true,
            quit: false,
            clipboard: ClipboardContext::new().unwrap(),
            action: String::new(),
            stashed_input: Vec::new(),
        }
    }

    pub async fn run(&mut self, mut terminal: Terminal<CrosstermBackend<Stdout>>) -> Result<()> {
        let mut input_events = EventStream::new();
        let mut redraw_tick = interval(Duration::from_millis(REDRAW_TICK));
        let mut activity_refresh_counter = 0;

        // Main loop
        while !self.quit {
            tokio::select! {
                    Some(Ok(event)) = input_events.next() => {
                        self.handle_input_event(event).await;
                    }

                    Some(event) = self.rx.recv() => {
                        self.handle_session_event(event);
                    }

                    _ = redraw_tick.tick() => {
                        // @todo catch errors instead of panicing to safe history
                        activity_refresh_counter = self.update_activity_indication(activity_refresh_counter);

                        if self.dirty {
                            terminal.draw(|frame| self.render(frame))?;
                            self.dirty = false;
                        }
                    }
            }
        }

        Ok(())
    }

    fn render(&mut self, frame: &mut Frame) {
        let [content, status_area] =
            Layout::vertical([Constraint::Min(0), Constraint::Length(1)]).areas(frame.area());

        let input_height = 3;

        let [conversation_area, input_area] =
            Layout::vertical([Constraint::Min(0), Constraint::Length(input_height)]).areas(content);

        let [input_indicator_area, input_field_area] =
            Layout::horizontal([Constraint::Length(2), Constraint::Min(10)]).areas(input_area);

        //How can I inspect the raw provider/rig response to debug tool call issues?
        frame.render_stateful_widget(ConversationView, conversation_area, &mut self.conversation);
        frame.render_widget("❯", input_indicator_area); // , , >,, ⟩, », 》, ▸ or ▶, ❯[  ▶︎ ] 
        frame.render_widget(&self.input, input_field_area);

        // update lines after redraw
        self.status.current_line = self.conversation.current_line();
        self.status.total_lines = self.conversation.total_lines();
        // @todo cloning is ugly, either build a stateful widget or a component
        frame.render_widget(self.status.clone(), status_area);
    }

    async fn handle_input_event(&mut self, event: Event) {
        if matches!(event, Event::Resize(_, _)) {
            self.dirty = true;
        } else {
            match self.status.input_mode {
                InputMode::Normal => self.handle_input_event_normal_mode(event).await,
                InputMode::Insert => self.handle_input_event_insert_mode(event).await,
                InputMode::View => self.handle_input_event_view_mode(event).await,
                InputMode::Command => self.handle_input_event_command_mode(event).await,
            }
        }
    }

    async fn handle_input_event_command_mode(&mut self, event: Event) {
        if let Event::Key(key) = event {
            match key.code {
                KeyCode::Esc => {
                    self.restore_input();
                    self.status.input_mode = InputMode::Normal;
                }
                KeyCode::Left => {
                    self.input.move_cursor(CursorMove::Back);
                }
                KeyCode::Right => {
                    self.input.move_cursor(CursorMove::Forward);
                }
                KeyCode::Enter => {
                    if let Some(command) = self.input.lines().first() {
                        self.session_tx
                            .send(SessionMessage::SubmitCommand(command.to_string()))
                            .await
                            .unwrap();
                    }
                    self.restore_input();
                    self.status.input_mode = InputMode::Normal;
                }
                _ => {
                    self.input.input_without_shortcuts(key);
                    // @todo ensure that first character always is a colon
                }
            }
            self.dirty = true;
        }
    }

    fn restore_input(&mut self) {
        self.input.clear();
        let mut has_previous_line = false;
        for line in &self.stashed_input {
            if has_previous_line {
                self.input.insert_newline();
            }
            self.input.insert_str(line);
            has_previous_line = true;
        }
        self.stashed_input.clear();
        // In restore_input() or when exiting command mode
        // self.input.skip(0); // or self.input.set_skip(0)
    }

    async fn handle_input_event_insert_mode(&mut self, event: Event) {
        if let Event::Key(key) = event {
            match key.code {
                KeyCode::Esc => {
                    self.status.input_mode = InputMode::Normal;
                }
                KeyCode::Up => {
                    self.input.move_cursor(CursorMove::Up);
                }
                KeyCode::Down => {
                    self.input.move_cursor(CursorMove::Down);
                }
                KeyCode::Left => {
                    self.input.move_cursor(CursorMove::Back);
                }
                KeyCode::Right => {
                    self.input.move_cursor(CursorMove::Forward);
                }
                KeyCode::PageUp => {
                    self.conversation.move_page_up();
                }
                KeyCode::PageDown => {
                    self.conversation.move_page_down();
                }
                KeyCode::Enter if key.modifiers == KeyModifiers::CONTROL => {
                    if let Some(prompt) = self.input.lines().first() {
                        self.session_tx
                            .send(SessionMessage::SubmitPrompt(prompt.to_string()))
                            .await
                            .unwrap();
                        self.input.clear();
                    }
                    self.status.input_mode = InputMode::Normal;
                }
                _ => {
                    self.input.input_without_shortcuts(key);
                }
            }
            self.dirty = true;
        }
    }

    async fn handle_input_event_normal_mode(&mut self, event: Event) {
        if let Event::Key(key) = event
            && (key.modifiers.is_empty() || key.modifiers == KeyModifiers::SHIFT)
        {
            match key.code {
                KeyCode::Esc => {
                    // maybe show a message and allow double esc for exit
                }
                KeyCode::Up => {
                    self.conversation.move_line_up();
                }
                KeyCode::Down => {
                    self.conversation.move_line_down();
                }
                KeyCode::PageUp => {
                    self.conversation.move_page_up();
                }
                KeyCode::PageDown => {
                    self.conversation.move_page_down();
                }
                KeyCode::Enter => {
                    if let Some(prompt) = self.input.lines().first() {
                        self.session_tx
                            .send(SessionMessage::SubmitPrompt(prompt.to_string()))
                            .await
                            .unwrap();
                        self.input.clear();
                    }
                }
                KeyCode::Char(c) => {
                    self.action.push(c);
                    self.resolve_action();
                }
                // KeyCode::Backspace => todo!(),
                // KeyCode::Home => todo!(),
                // KeyCode::End => todo!(),
                // KeyCode::Tab => todo!(),
                // KeyCode::Delete => todo!(),
                _ => {}
            }
            self.dirty = true;
        }
    }

    async fn handle_input_event_view_mode(&mut self, event: Event) {
        if let Event::Key(key) = event {
            match key.code {
                KeyCode::Esc => {
                    self.status.input_mode = InputMode::Normal;
                }
                KeyCode::Up => {
                    self.conversation.move_line_up();
                }
                KeyCode::Down => {
                    self.conversation.move_line_down();
                }
                KeyCode::PageUp => {
                    self.conversation.move_page_up();
                }
                KeyCode::PageDown => {
                    self.conversation.move_page_down();
                }
                KeyCode::Home => { // @todo go to start
                }
                KeyCode::End => { // @todo go to end
                }
                KeyCode::Tab => { // @todo go to next element or previous element (with shift)
                }
                _ => {}
            }
            self.dirty = true;
        }
    }

    fn handle_session_event(&mut self, event: UiMessage) {
        match event {
            UiMessage::ActiveModel(model) => {
                self.status.model = model
                    .and_then(|m| m.split('/').next_back().map(|s| s.to_string()))
                    .unwrap_or("No model selected".to_string());
            }
            UiMessage::ContextState { fill_state, total } => {
                self.status.context_used = fill_state;
                self.status.context_total = total;
            }
            UiMessage::ConversationItem(item) => {
                self.conversation.upsert_item(item);
                self.dirty = true;
            }
            UiMessage::ConversationViewMode(view_mode) => {
                self.conversation.set_view_mode(view_mode);
                self.dirty = true;
            }
            UiMessage::GenerationFinished(usage) => {
                self.status.tokens = usage;
            }
            UiMessage::IsProcessing(active) => {
                self.processing = active;
            }
            UiMessage::Quit => {
                self.quit = true;
            }
        }
    }

    fn resolve_action(&mut self) {
        if let Some(modal_action) = self.parse_action() {
            match modal_action {
                ModalAction::CommandMode(c) => {
                    self.stashed_input = self.input.lines().to_vec();
                    self.input.clear();
                    self.status.input_mode = InputMode::Command;
                    self.input.insert_str(c.to_string());
                    // Make the first character (':') non-editable
                    //self.input.skip(1);
                }
                ModalAction::Goto(_position) => {
                    // @todo implement
                }
                ModalAction::InsertMode => {
                    self.status.input_mode = InputMode::Insert;
                }
                ModalAction::Invalid(_action) => {
                    // @todo report error
                }
                ModalAction::MoveDown => {
                    self.input.move_cursor(CursorMove::Down);
                }
                ModalAction::MoveBack => {
                    self.input.move_cursor(CursorMove::Forward);
                }
                ModalAction::MoveForward => {
                    self.input.move_cursor(CursorMove::Back);
                }
                ModalAction::MoveUp => {
                    self.input.move_cursor(CursorMove::Up);
                }
                ModalAction::MoveWordBack => {
                    self.input.move_cursor(CursorMove::WordBack);
                }
                ModalAction::MoveWordEnd => {
                    self.input.move_cursor(CursorMove::WordEnd);
                }
                ModalAction::MoveWordForward => {
                    self.input.move_cursor(CursorMove::WordForward);
                }
                ModalAction::Paste => {
                    // @todo error handling/reporting
                    let content = self.clipboard.get_contents().unwrap_or_default();
                    self.input.insert_str(content);
                }
                ModalAction::Redo => {
                    self.input.redo();
                }
                ModalAction::Undo => {
                    self.input.undo();
                }
                ModalAction::ViewMode => {
                    self.status.input_mode = InputMode::View;
                }
                ModalAction::YankSelection => {
                    // @todo implement
                }
                ModalAction::YankSnippet(position) => {
                    match self.find_code_snippet(position) {
                        Some(code_snippet) => {
                            // @todo report error instead of unwrap
                            self.clipboard
                                .set_contents(code_snippet.to_owned())
                                .unwrap();
                        }
                        None => {
                            // @todo report error: snipped NR not found
                        }
                    }
                }
            }
            self.action.clear();
        }
    }

    fn parse_action(&mut self) -> Option<ModalAction> {
        if self.action.is_empty() {
            return None;
        }

        match self.action.as_str() {
            "b" => Some(ModalAction::MoveWordBack),
            "e" => Some(ModalAction::MoveWordEnd),
            "i" => Some(ModalAction::InsertMode),
            "h" => Some(ModalAction::MoveForward),
            "j" => Some(ModalAction::MoveDown),
            "k" => Some(ModalAction::MoveUp),
            "l" => Some(ModalAction::MoveBack),
            "p" => Some(ModalAction::Paste),
            "u" => Some(ModalAction::Undo),
            "U" => Some(ModalAction::Redo),
            "v" => Some(ModalAction::ViewMode),
            "w" => Some(ModalAction::MoveWordForward),
            "y" => Some(ModalAction::YankSelection),
            ":" => Some(ModalAction::CommandMode(':')),
            _ => {
                // Try parsing as <number><command>
                // operate on chars to ensure slicing at character boundaries only
                let chars: Vec<char> = self.action.chars().collect();
                if chars.len() < 2 {
                    return None;
                }
                if chars.iter().all(|c| c.is_numeric()) {
                    return None;
                }

                let last_char = chars.last()?;
                let num_part: String = chars[..chars.len() - 1].iter().collect();

                // Ensure prefix is purely numeric and suffix is not
                if let Ok(position) = num_part.parse::<u32>()
                    && !last_char.is_numeric()
                {
                    match last_char {
                        'y' => Some(ModalAction::YankSnippet(position)),
                        'g' => Some(ModalAction::Goto(position)),
                        _ => Some(ModalAction::Invalid(self.action.clone())),
                    }
                } else {
                    Some(ModalAction::Invalid(self.action.clone()))
                }
            }
        }
    }

    fn update_activity_indication(&mut self, counter: u64) -> u64 {
        if counter == 0 {
            if self.processing {
                self.status.indicate_activity();
            } else {
                self.status.indicate_idle();
            }
            self.dirty = true;

            ACTIVITY_REFRESH_INTERVAL
        } else {
            counter.saturating_sub(1)
        }
    }

    fn find_code_snippet(&self, copy_id: u32) -> Option<String> {
        let mut snippet = None;
        let mut backlog = Vec::new();

        for item in &self.conversation.conversation_items {
            match &item.content {
                ConversationItemKind::UserMessage(message_content) => {
                    backlog.extend(message_content.content().blocks.clone());
                }
                ConversationItemKind::AssistantMessage(message_content, _) => {
                    backlog.extend(message_content.content().blocks.clone());
                }
                ConversationItemKind::AssistantReasoning(message_content) => {
                    backlog.extend(message_content.content().blocks.clone());
                }
                ConversationItemKind::ToolInvocation(_) => {}
                ConversationItemKind::ToolOutput(_) => {}
                ConversationItemKind::Notification(_, _) => {}
            }
        }

        while let Some(current_block) = backlog.pop()
            && snippet.is_none()
        {
            match current_block {
                Block::Code {
                    lang: _lang,
                    content,
                    copy_id: candidate_copy_id,
                } if copy_id == candidate_copy_id => snippet = Some(content),
                Block::List { content } => {
                    for item in content {
                        backlog.extend(item.content.clone());
                    }
                }
                Block::Quote(blocks) => {
                    backlog.extend(blocks.clone());
                }
                _ => {}
            }
        }

        snippet
    }
}