monothink 0.5.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 monothink.
 * If not, see <https://www.gnu.org/licenses/>.
 */

use ratatui::{
    buffer::Buffer,
    layout::Rect,
    style::{Color, Modifier, Style},
    text::{Line, Span},
    widgets::{Paragraph, StatefulWidget, Widget},
};
use rig::message::{CallId, LocalCallId};
use std::collections::HashMap;
use unicode_width::UnicodeWidthChar;

use crate::{
    common::conversation::{
        ConversationItem, ConversationItemKind, Document, Inline, MessageState, NotificationType,
        TableCell, TableRow, ToolInvocationContent, ToolOutputContent,
    },
    config::ViewMode,
    ui::{
        md_renderer::{RenderConfig, add_gutter, render_block, render_table},
        style::{
            DARK_BLUE, DARK_PURPLE, GUTTER_WIDTH, LIGHT_BLUE, LIGHT_TOURQUOISE, PURPLE, TOURQUOISE,
            TextStyle,
        },
    },
};

// not really a page but a convenient step size
const PAGE_SIZE: u64 = 40;

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum ScrollingPosition {
    #[default]
    Auto,
    Fixed(u64),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Movement {
    Start,
    End,
    Line(u32),
    LineUp,
    LineDown,
    PageUp,
    PageDown,
}

#[derive(Clone, Debug, Default)]
pub struct SymbolConfig {
    pub assistant_message: String,
    pub assistant_reasoning: String,
    pub tool_message: String,
    pub message_state_final: String,
    pub message_state_intermediate: String,
    pub message_notification_debug: String,
    pub message_notification_error: String,
    pub message_notification_info: String,
    pub message_notification_warning: String,
    pub status_failed: String,
    pub status_success: String,
    pub system_message: String,
    pub user_message: String,
}

// #[derive(Clone, Debug)]
// pub struct ItemView {
//     pub item: ConversationItem,
//     pub view: Vec<Line<'static>>,
// }

#[derive(Clone, Debug, Default)]
pub struct ConversationViewModel {
    conversation_view: Vec<Line<'static>>,
    // @todo instead of this, have a ItemView for each item, containing both the
    // item and its rendered line. When a item is updated or added,
    // lines is emtpy. On rendering, the widget goes through all items and renders
    // those that are not rendered. Rendering then becomes a function of ItemView
    // and not the widget.
    // Finally, the lines are assembled into one text (or later replaced by a custom widget)
    pub conversation_items: Vec<ConversationItem>,
    dirty: bool,
    height: Option<u64>,
    scrolling_position: ScrollingPosition,
    render_config: RenderConfig,
    symbol_cfg: SymbolConfig,
    view_mode: ViewMode,
}

#[derive(Debug, Clone, Default)]
pub struct ConversationView;

impl StatefulWidget for ConversationView {
    type State = ConversationViewModel;

    fn render(self, area: Rect, buf: &mut Buffer, state: &mut Self::State) {
        state.height = Some(area.height as u64);
        let height = state
            .height
            .expect("Height must be set after assignment above");

        if state.dirty {
            let config = RenderConfig {
                width: area.width,
                ..state.render_config.clone()
            };
            state.conversation_view = Self::render_conversation(
                &state.conversation_items,
                &config,
                &state.symbol_cfg,
                state.view_mode,
            );

            // @todo red rectangle when clipped to show an error in debug build
            Self::clip_lines(config.width, &mut state.conversation_view);

            state.dirty = false;
        }

        let total_lines = state.conversation_view.len();
        let scroll_offset = state.scroll_offset(height, total_lines);

        if total_lines == 0 {
            return;
        }

        let viewport_start = scroll_offset as usize;
        let viewport_end = std::cmp::min(viewport_start + height as usize, total_lines);
        let visible_lines = state.conversation_view[viewport_start..viewport_end].to_vec();

        Paragraph::new(visible_lines)
            .style(TextStyle::default())
            .render(area, buf);
    }
}

impl ConversationViewModel {
    pub fn new(render_config: RenderConfig, symbol_cfg: SymbolConfig) -> Self {
        Self {
            render_config,
            symbol_cfg,
            ..Default::default()
        }
    }

    pub fn clear(&mut self) {
        self.conversation_view.clear();
        self.conversation_items.clear();
        self.dirty = true;
    }

    pub fn current_line(&self) -> u64 {
        if let Some(height) = self.height {
            self.scroll_offset(height, self.conversation_view.len())
        } else {
            0
        }
    }

    pub fn total_lines(&self) -> u64 {
        self.conversation_view.len() as u64
    }

    pub fn upsert_item(&mut self, item: ConversationItem) {
        self.dirty = true;
        let maybe_outdated_item = self
            .conversation_items
            .iter_mut()
            .rev()
            .find(|it| it.id == item.id);
        if let Some(outdated_item) = maybe_outdated_item {
            outdated_item.content = item.content;
        } else {
            self.conversation_items.push(item);
        }
    }

    pub fn move_to(&mut self, movement: Movement) {
        match movement {
            Movement::Start => self.scrolling_position = ScrollingPosition::Fixed(0),
            Movement::End => self.scrolling_position = ScrollingPosition::Auto,
            Movement::Line(_) => todo!(),
            Movement::LineUp => self.move_up(1),
            Movement::LineDown => self.move_down(1),
            Movement::PageUp => self.move_up(PAGE_SIZE),
            Movement::PageDown => self.move_down(PAGE_SIZE),
        }
    }

    pub fn set_view_mode(&mut self, view_mode: ViewMode) {
        self.dirty = true;
        self.view_mode = view_mode;
    }

    fn move_down(&mut self, step: u64) {
        if !self.can_scroll() {
            return;
        }

        let total_lines = self.conversation_view.len() as u64;
        if let Some(height) = self.height {
            let max_offset = total_lines.saturating_sub(height);
            let new_offset = self
                .scroll_offset(height, total_lines as usize)
                .saturating_add(step);
            if new_offset < max_offset {
                self.scrolling_position = ScrollingPosition::Fixed(new_offset);
            } else {
                self.scrolling_position = ScrollingPosition::Auto;
            }
        }
    }

    fn move_up(&mut self, step: u64) {
        let p = self.scroll_offset(self.height.unwrap_or(0), self.conversation_view.len());
        if p >= step {
            self.scrolling_position = ScrollingPosition::Fixed(p - step);
        } else {
            self.scrolling_position = ScrollingPosition::Fixed(0);
        }
    }

    fn can_scroll(&self) -> bool {
        if let Some(height) = self.height {
            (height as usize) < self.conversation_view.len()
        } else {
            false
        }
    }

    /// Returns the scroll offset (first visible line index) given the viewport height
    /// and the total number of lines.
    fn scroll_offset(&self, viewport_height: u64, total_lines: usize) -> u64 {
        let total_lines = total_lines as u64;

        if total_lines <= viewport_height || viewport_height == 0 {
            return 0;
        }

        match self.scrolling_position {
            ScrollingPosition::Auto => {
                // Auto-scroll: show the bottom of the conversation
                total_lines.saturating_sub(viewport_height)
            }
            ScrollingPosition::Fixed(offset) => {
                // Clamp to valid range
                std::cmp::min(offset, total_lines.saturating_sub(viewport_height))
            }
        }
    }
}

impl ConversationView {
    pub fn render_conversation(
        items: &[ConversationItem],
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
        view_mode: ViewMode,
    ) -> Vec<Line<'static>> {
        match view_mode {
            ViewMode::Debug => items
                .iter()
                .flat_map(|item| Self::render_item(item, render_config, symbols_cfg))
                .collect(),
            ViewMode::Full => items
                .iter()
                .filter(|item| !item.is_notification(Some(NotificationType::Debug)))
                .flat_map(|item| Self::render_item(item, render_config, symbols_cfg))
                .collect(),
            ViewMode::Detailed => {
                let mut tool_calls = HashMap::new();
                let mut tool_call_pairs = Vec::new();
                items
                    .iter()
                    .filter(|item| !item.is_notification(Some(NotificationType::Debug)))
                    .flat_map(|item| match &item.content {
                        ConversationItemKind::ToolInvocation(tool_invocation_content) => {
                            tool_calls.insert(
                                tool_invocation_content.call_id.clone(),
                                tool_invocation_content,
                            );
                            Vec::new()
                        }
                        ConversationItemKind::ToolOutput(tool_output_content) => {
                            let tool_call = tool_calls
                                .get(&tool_output_content.call_id)
                                .map(|tc| tc.to_owned().clone())
                                .unwrap_or_else(|| ToolInvocationContent {
                                    name: "N/A".to_string(),
                                    arguments: serde_json::Map::new(),
                                    call_id: CallId::Local(LocalCallId::new()),
                                });
                            tool_call_pairs.push((tool_call, tool_output_content.clone()));
                            Vec::new()
                        }
                        _ => {
                            let mut lines = Vec::new();
                            if !tool_call_pairs.is_empty() {
                                lines.extend(Self::render_tool_invocations(
                                    &tool_call_pairs,
                                    render_config,
                                    symbols_cfg,
                                ));
                                lines.push(Line::raw(" "));
                                tool_call_pairs.clear();
                            }
                            lines.extend(Self::render_item(item, render_config, symbols_cfg));
                            lines
                        }
                    })
                    .collect()
            }
            ViewMode::Compact => {
                let mut output = Vec::new();
                for interaction in Self::interactions(items) {
                    let (user_messages, mut tool_invocations, mut notifications) =
                        Self::render_interaction_parts(&interaction, render_config, symbols_cfg);
                    let assistant_output = Self::render_final_or_intermediate_assistant(
                        render_config,
                        symbols_cfg,
                        &interaction,
                    );
                    Self::conditionally_add_separator(&mut tool_invocations);
                    Self::conditionally_add_separator(&mut notifications);

                    output.extend(user_messages);
                    output.extend(tool_invocations);
                    output.extend(assistant_output);
                    output.extend(notifications);
                }

                output
            }
            ViewMode::Minimal => {
                let mut output = Vec::new();
                for interaction in Self::interactions(items) {
                    let (user_messages, _tool_invocations, notifications) =
                        Self::render_interaction_parts(&interaction, render_config, symbols_cfg);
                    let assistant_output = Self::render_final_or_intermediate_assistant(
                        render_config,
                        symbols_cfg,
                        &interaction,
                    );
                    output.extend(user_messages);
                    output.extend(assistant_output);
                    output.extend(notifications);
                }

                output
            }
        }
    }

    // renders the item to lines
    // For streamable items (esp. AgentMessage) it renders the raw content
    // as preliminary output, if no structured content is available.
    // we should have a proper type for the render state we need
    pub fn render_item(
        item: &ConversationItem,
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        let mut lines = match &item.content {
            ConversationItemKind::UserMessage(msg) => {
                Self::render_user(msg.content(), render_config, symbols_cfg)
            }
            ConversationItemKind::AssistantMessage(msg, kind) => {
                Self::render_assistant(msg.content(), render_config, symbols_cfg, kind)
            }
            ConversationItemKind::AssistantReasoning(msg) => {
                Self::render_reasoning(msg.content(), render_config, symbols_cfg)
            }
            ConversationItemKind::Notification(notification_type, content) => {
                Self::render_notification(
                    notification_type.to_owned(),
                    content,
                    render_config.width,
                    symbols_cfg,
                )
            }
            ConversationItemKind::ToolInvocation(ToolInvocationContent {
                name,
                arguments,
                call_id,
            }) => Self::render_tool_invocation(name, arguments, call_id, symbols_cfg),
            ConversationItemKind::ToolOutput(ToolOutputContent {
                name,
                output,
                success,
                call_id,
            }) => Self::render_tool_result(
                name,
                output,
                *success,
                call_id,
                render_config.width,
                symbols_cfg,
            ),
        };

        // spacing between items
        lines.push(Line::raw(""));

        lines
    }

    fn render_assistant(
        msg: &Document,
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
        kind: &MessageState,
    ) -> Vec<Line<'static>> {
        let kind_icon = match kind {
            MessageState::Intermediate => symbols_cfg.message_state_intermediate.clone(),
            MessageState::Final => symbols_cfg.message_state_final.clone(),
        };
        Self::render_message(
            |w| {
                let mut lines = Vec::new();
                let config = RenderConfig {
                    highlight_cfg: render_config.highlight_cfg.clone(),
                    style: render_config.style.clone(),
                    width: w,
                };
                for block in &msg.blocks {
                    lines.extend(render_block(block, &config));
                }
                lines
            },
            LIGHT_BLUE,
            DARK_BLUE,
            render_config.width,
            &format!(
                "{} Assistant ({})",
                symbols_cfg.assistant_message, kind_icon
            ),
        )
    }

    fn render_reasoning(
        msg: &Document,
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        Self::render_message(
            |w| {
                let mut lines = Vec::new();
                let config = RenderConfig {
                    highlight_cfg: render_config.highlight_cfg.clone(),
                    style: render_config.style.clone(),
                    width: w,
                };
                for block in &msg.blocks {
                    lines.extend(render_block(block, &config));
                }
                lines
            },
            LIGHT_TOURQUOISE,
            TOURQUOISE,
            render_config.width,
            &format!("{} Reasoning", symbols_cfg.assistant_reasoning),
        )
    }

    fn render_user(
        msg: &Document,
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        Self::render_message(
            |w| {
                let mut lines = Vec::new();
                let config = RenderConfig {
                    highlight_cfg: render_config.highlight_cfg.clone(),
                    style: render_config.style.clone(),
                    width: w,
                };
                for block in &msg.blocks {
                    lines.extend(render_block(block, &config));
                }
                lines
            },
            PURPLE,
            DARK_PURPLE,
            render_config.width,
            &format!("{} You", symbols_cfg.user_message),
        )
    }

    fn render_final_or_intermediate_assistant(
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
        interaction: &[ConversationItem],
    ) -> Vec<Line<'static>> {
        let assistant_output: Vec<_> = if interaction
            .iter()
            .any(|item| item.is_assistant_message(Some(MessageState::Final)))
        {
            interaction
                .iter()
                .filter(|msg| msg.is_assistant_message(Some(MessageState::Final)))
                .flat_map(|item| Self::render_item(item, render_config, symbols_cfg))
                .collect()
        } else {
            interaction
                .iter()
                .filter(|msg| {
                    if let ConversationItemKind::AssistantMessage(msg_content, _) = &msg.content {
                        !msg_content.raw().trim().is_empty()
                    } else {
                        false
                    }
                })
                .flat_map(|item| Self::render_item(item, render_config, symbols_cfg))
                .collect()
        };
        assistant_output
    }

    fn render_interaction_parts(
        interaction: &[ConversationItem],
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
    ) -> (Vec<Line<'static>>, Vec<Line<'static>>, Vec<Line<'static>>) {
        let user_messages = interaction
            .iter()
            .filter(|item| item.is_user_message())
            .flat_map(|item| Self::render_item(item, render_config, symbols_cfg))
            .collect();
        let tool_invocations = Self::render_tool_invocations(
            &tool_invocations(interaction),
            render_config,
            symbols_cfg,
        );
        let notifications = Self::render_notifications(
            interaction.iter().filter(|msg| {
                msg.is_notification(None) && !msg.is_notification(Some(NotificationType::Debug))
            }),
            symbols_cfg,
        );
        (user_messages, tool_invocations, notifications)
    }

    fn render_message<T: FnOnce(u16) -> Vec<Line<'static>>>(
        body: T,
        text_color: Color,
        gutter_color: Color,
        width: u16,
        header: &str,
    ) -> Vec<Line<'static>> {
        let mut lines = vec![Self::render_header(header, text_color), Line::raw(" ")];

        lines.extend(body(width.saturating_sub(*GUTTER_WIDTH)));

        add_gutter(lines, gutter_color, None)
    }
    // @todo take an iterator and yield another iterator, avoiding building all those vectors
    fn clip_lines(width: u16, lines: &mut Vec<Line<'_>>) {
        for line in lines {
            if line.width() > width as usize {
                // @todo issue a tracing warning, because my line calculation should avoid this scenario
                let mut total_width = 0;
                let mut new_spans = Vec::new();
                for span in &mut line.spans {
                    let span_width = span.width();
                    if total_width + span_width <= width as usize {
                        total_width += span_width;
                        new_spans.push(span.clone());
                    } else if total_width < width as usize {
                        let new_text: String = span
                            .content
                            .chars()
                            .take_while(|c| {
                                total_width += c.width().unwrap_or_default();
                                total_width <= width as usize
                            })
                            .collect();
                        span.content = new_text.into();
                        new_spans.push(span.clone());
                        break;
                    } else {
                        break;
                    }
                }
                line.spans = new_spans;
            }
        }
    }

    fn conditionally_add_separator(lines: &mut Vec<Line<'_>>) {
        if !lines.is_empty() {
            lines.push(Line::raw(" "));
        }
    }

    fn render_header(label: &str, color: Color) -> Line<'static> {
        Line::from(vec![Span::styled(
            label.to_string(),
            Style::default().fg(color).add_modifier(Modifier::BOLD),
        )])
    }

    fn render_tool_invocation(
        name: &str,
        args: &serde_json::Map<String, serde_json::Value>,
        _call_id: &CallId,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        let lines = vec![
            Self::render_header(
                &format!("{} Tool: {}", symbols_cfg.tool_message, name),
                Color::Rgb(220, 220, 170),
            ),
            Line::raw(" "),
            Line::from(Span::styled(
                format!("> {}", serde_json::Value::Object(args.clone())),
                Style::default().fg(Color::DarkGray),
            )),
        ];

        add_gutter(lines, Color::Yellow, None)
    }

    fn render_tool_invocations(
        invocations: &[(ToolInvocationContent, ToolOutputContent)],
        render_config: &RenderConfig,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        if !invocations.is_empty() {
            let mut lines = vec![
                Self::render_header(
                    &format!("{} Tools", symbols_cfg.tool_message),
                    Color::Rgb(220, 220, 170),
                ),
                Line::raw(" "),
            ];
            let mut rows = vec![TableRow {
                content: vec![
                    TableCell {
                        content: vec![Inline::Text("Invocation".to_string())],
                    },
                    TableCell {
                        content: vec![Inline::Text("Result".to_string())],
                    },
                ],
            }];

            for (invocation, result) in invocations {
                let success_icon = if result.success {
                    &symbols_cfg.status_success
                } else {
                    &symbols_cfg.status_failed
                };
                rows.push(TableRow {
                    content: vec![
                        TableCell {
                            content: vec![Inline::Text(format!(
                                " {}: {}",
                                invocation.name,
                                serde_json::Value::Object(invocation.arguments.clone())
                            ))],
                        },
                        TableCell {
                            content: vec![Inline::Text(success_icon.to_string())],
                        },
                    ],
                });
            }

            lines.extend(render_table(&rows, render_config));
            add_gutter(lines, Color::Yellow, None)
        } else {
            Vec::new()
        }
    }

    fn render_tool_result(
        name: &str,
        output: &str,
        success: bool,
        _call_id: &CallId,
        width: u16,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        let mut lines = vec![
            Self::render_header(
                &format!("{} Tool Output", symbols_cfg.tool_message),
                Color::Rgb(220, 220, 170),
            ),
            Line::raw(" "),
        ];

        for line in output.lines() {
            lines.push(Line::from(Span::raw(line.to_string())));
        }

        add_gutter(lines, Color::Yellow, None)
    }

    fn render_notification(
        notification_type: NotificationType,
        msg: &str,
        width: u16,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>> {
        let lines = vec![
            Self::render_header(
                &format!("{} System", symbols_cfg.system_message),
                Color::Green,
            ),
            Line::from(vec![
                Span::raw(format!(
                    "{} ",
                    notification_icon(notification_type, symbols_cfg)
                )),
                Span::raw(msg.to_string()),
            ]),
        ];

        add_gutter(lines, Color::Rgb(62, 143, 126), None)
    }

    fn render_notifications<'a, I>(
        notifications: I,
        symbols_cfg: &SymbolConfig,
    ) -> Vec<Line<'static>>
    where
        I: Iterator<Item = &'a ConversationItem>,
    {
        let mut lines = vec![Self::render_header(
            &format!("{} System", symbols_cfg.system_message),
            Color::Green,
        )];
        let mut has_notifications = false;
        for notification in notifications {
            if let ConversationItemKind::Notification(notification_type, msg) =
                &notification.content
            {
                has_notifications = true;
                let spans = vec![
                    Span::raw(format!(
                        "{} ",
                        notification_icon(*notification_type, symbols_cfg)
                    )),
                    Span::raw(msg.to_string()),
                ];
                lines.push(Line::from(spans));
            }
        }

        if has_notifications {
            add_gutter(lines, Color::Rgb(62, 143, 126), None)
        } else {
            Vec::new()
        }
    }

    // I am surprised there isn't a generic iterator for this.
    // split_inclusive from Vec comes close but includes the separator as a terminator.
    // There does not seem to be a function that includes it as a start element.
    // This would be more efficient as a iterator but so far, this simple version is sufficient
    fn interactions(items: &[ConversationItem]) -> Vec<Vec<ConversationItem>> {
        let mut result = Vec::new();
        if items.is_empty() {
            return result;
        }
        let mut chunk = vec![items.first().unwrap().to_owned()];
        for item in items.iter().skip(1) {
            if item.is_user_message() {
                result.push(chunk);
                chunk = vec![item.to_owned()];
            } else {
                chunk.push(item.to_owned());
            }
        }
        if !chunk.is_empty() {
            result.push(chunk);
        }
        result
    }
}

fn notification_icon(notification_type: NotificationType, symbols_cfg: &SymbolConfig) -> String {
    match notification_type {
        NotificationType::Debug => symbols_cfg.message_notification_debug.clone(),
        NotificationType::Info => symbols_cfg.message_notification_info.clone(),
        NotificationType::Warn => symbols_cfg.message_notification_warning.clone(),
        NotificationType::Error => symbols_cfg.message_notification_error.clone(),
    }
}

fn tool_invocations(
    interaction: &[ConversationItem],
) -> Vec<(ToolInvocationContent, ToolOutputContent)> {
    let tool_results: Vec<_> = interaction
        .iter()
        .filter_map(|item| {
            if let ConversationItem {
                content: ConversationItemKind::ToolOutput(output),
                id: _,
            } = item
            {
                Some(output)
            } else {
                None
            }
        })
        .collect();

    interaction
        .iter()
        .filter_map(|item| {
            if let ConversationItem {
                content: ConversationItemKind::ToolInvocation(invocation),
                id: _,
            } = item
            {
                tool_results
                    .iter()
                    .find(|res| res.call_id == invocation.call_id)
                    .map(|result| (invocation.to_owned(), (*result).clone()))
            } else {
                None
            }
        })
        .collect()
}