monothink 0.5.1

/*
 * 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 std::{
    env, fs,
    io::Write,
    path::{Path, PathBuf},
};

use anyhow::{Error, Result};
use directories::ProjectDirs;
use rig::{completion::Usage, message::CallId};
use tokio::sync::mpsc::{Receiver, Sender};

use crate::{
    common::conversation::{
        Conversation, ConversationItemKind, Inline, MessageState, NotificationType,
    },
    config::{Config, Endpoint, ViewMode, default_context_window},
    llm::LlmMessage,
    ui::{UiMessage, suggestion_popup::Suggestion},
};

#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct TokenUsage {
    pub input: Option<u64>,
    pub output: Option<u64>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Command {
    CancelRequest,
    Export(String),
    Quit,
    SelectModel(String),
    Title(String),
    ViewMode(ViewMode),
}

#[derive(Clone, Debug, PartialEq)]
pub enum SessionMessage {
    SubmitCommand(Command),
    SubmitPrompt(String),
    AgentResponse(String),
    AgentReasoning(String),
    AgentCompletion(Usage), // Usage is a rig type we we should abstract from. But at this point we need the same information anyhow.
    ContextPressure {
        initial: u64,
        updated: u64,
        window: u64,
    },
    Debug(String),
    FinalAgentResponse(String),
    Timeout,
    ToolCall {
        name: String,
        args: serde_json::Value,
        call_id: CallId,
    },
    ToolResult {
        name: String,
        output: String,
        success: bool,
        call_id: CallId,
    },
    Error(String),
    Info(String),
    Warning(String),
}

#[derive(Debug)]
pub struct Session {
    config: Config,
    active_model: Option<String>,
    ui_tx: Sender<UiMessage>,
    llm_tx: Sender<LlmMessage>,
    rx: Receiver<SessionMessage>,
    context_window: u64,
    conversation: Conversation,
    token_usage: TokenUsage,
    quit: bool,
}

impl Session {
    pub fn new(
        ui_tx: Sender<UiMessage>,
        llm_tx: Sender<LlmMessage>,
        rx: Receiver<SessionMessage>,
        config: Config,
        workspace_root: PathBuf,
    ) -> Self {
        let endpoint = Self::get_primary_endpoint(&config);
        let active_model = match endpoint {
            Ok(endpoint) => Some(endpoint.default_model),
            Err(_) => None,
        };

        let context_window = match &active_model {
            Some(model) => context_window(&config, model),
            None => default_context_window(),
        };

        Self {
            config,
            active_model: active_model.clone(),
            ui_tx,
            llm_tx,
            rx,
            context_window,
            conversation: Conversation::new(active_model, workspace_root),
            token_usage: TokenUsage::default(),
            quit: false,
        }
    }

    pub async fn run(&mut self) {
        self.update_model_list().await;
        // @todo proper error handling: Tell user to configure an endpoint
        let endpoint = Self::get_primary_endpoint(&self.config).unwrap();
        self.llm_tx
            .send(LlmMessage::SetEndpoint(endpoint))
            .await
            .unwrap();
        if let Some(active_model) = &self.active_model {
            // @todo proper error handling: Tell user to configure a model
            self.set_model(active_model).await;
        }
        self.ui_tx
            .send(UiMessage::ConversationViewMode(
                self.config.general.view_mode,
            ))
            .await
            .unwrap();
        self.ui_tx
            .send(UiMessage::ContextState {
                fill_state: 0,
                total: self.context_window,
            })
            .await
            .unwrap();

        while !self.quit {
            let event = self.rx.recv().await;
            if let Some(msg) = event {
                match msg {
                    SessionMessage::SubmitCommand(command) => self.evaluate_command(command).await,
                    SessionMessage::SubmitPrompt(prompt) => {
                        let item = self.conversation.add_user_prompt(&prompt).clone();
                        self.llm_tx
                            .send(LlmMessage::SubmitPrompt {
                                prompt,
                                history: self.conversation.messages().clone(),
                            })
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.clone()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(true))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentResponse(chunk) => {
                        let item = self.conversation.add_assistant_response_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentReasoning(chunk) => {
                        let item = self.conversation.add_assistant_reasoning_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::AgentCompletion(usage) => {
                        self.update_token_usage(usage);
                        self.ui_tx
                            .send(UiMessage::GenerationFinished(TokenUsage {
                                input: self.token_usage.input,
                                output: self.token_usage.output,
                            }))
                            .await
                            .unwrap();

                        // publish report as system message for analyzing provider behavior
                        // bypass conversation as this is only for development
                        let msg = format!(
                            "Usage: input tokens = {}, output tokens = {},\n \
                            cached input tokens = {}, cached creation input tokens = {},\n \
                            reasoning tokens = {}, tool use prompt tokens = {},\n total tokens = {}",
                            usage.input_tokens.unwrap_or_default(),
                            usage.output_tokens.unwrap_or_default(),
                            usage.cached_input_tokens.unwrap_or_default(),
                            usage.cache_creation_input_tokens.unwrap_or_default(),
                            usage.reasoning_tokens.unwrap_or_default(),
                            usage.tool_use_prompt_tokens.unwrap_or_default(),
                            usage.total_tokens.unwrap_or_default()
                        );
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ContextPressure {
                        initial,
                        updated,
                        window,
                    } => {
                        let msg = format!(
                            "Context window: {window}, initial pressure: {initial}, updated pressure: {updated}.",
                        ).to_string();
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::ContextState {
                                fill_state: updated,
                                total: window,
                            })
                            .await
                            .unwrap();
                    }
                    SessionMessage::FinalAgentResponse(msg) => {
                        let item = self.conversation.set_final_assistant_response(&msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Error(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Error, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Warning(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Warn, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Timeout => {
                        let item = self.conversation.add_notification(
                            NotificationType::Error,
                            "The connection to the LLM provider timed out.".to_string(),
                        );
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                        self.ui_tx
                            .send(UiMessage::IsProcessing(false))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolCall {
                        name,
                        args,
                        call_id,
                    } => {
                        let item = self.conversation.add_tool_invocation(name, args, call_id);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::ToolResult {
                        name,
                        output,
                        success,
                        call_id,
                    } => {
                        let item = self
                            .conversation
                            .add_tool_result(name, output, success, call_id);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Debug(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Debug, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                    SessionMessage::Info(msg) => {
                        let item = self
                            .conversation
                            .add_notification(NotificationType::Info, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await
                            .unwrap();
                    }
                }
            }
        }
    }

    async fn evaluate_command(&mut self, command: Command) {
        match command {
            Command::CancelRequest => {
                // @todo handle cancelling, if possible
            }
            Command::Export(name) => self.handle_export_command(&name).await,
            Command::Quit => self.handle_quit_command().await,
            Command::SelectModel(model) => self.select_model(model).await,
            Command::Title(name) => {
                self.conversation.set_title(&name);
                self.ui_tx
                    .send(UiMessage::Title(
                        self.conversation.title().unwrap_or_default(),
                    ))
                    .await
                    .unwrap();
            }
            Command::ViewMode(view_mode) => self.handle_view_mode_command(view_mode).await,
        }
    }

    // @todo this should be part of config
    fn get_primary_endpoint(config: &Config) -> Result<Endpoint> {
        config
            .endpoint
            .get(&config.general.primary_endpoint)
            .map(|val| val.to_owned())
            .ok_or(Error::msg(format!(
                "Endpoint {} not found in endpoints.",
                config.general.primary_endpoint
            )))
    }

    async fn handle_export_command(&mut self, params: &str) {
        let mut path = Path::new(params).to_path_buf();
        // check if it is an absolute or relative path
        if Some(true) != path.parent().and_then(|p| p.try_exists().ok()) {
            let mut relative_path = env::current_dir().expect("Cannot get current directory");
            relative_path.push(params);
            path = relative_path;
        }

        let mut first = true;

        if Some(true) == path.parent().and_then(|p| p.try_exists().ok()) {
            // iterate over messages and write user and assistant messages as md
            let mut output = String::new();
            for msg in self.conversation.messages() {
                match &msg.content {
                    ConversationItemKind::UserMessage(message_content) => {
                        if !first {
                            output.push_str("\n\n");
                        } else {
                            first = false;
                        }
                        output.push_str("# User\n\n");
                        output.push_str(message_content.raw());
                    }
                    ConversationItemKind::AssistantMessage(message_content, message_kind) => {
                        if !message_content.raw().trim().is_empty()
                            && *message_kind == MessageState::Final
                        {
                            // note: Assistant can never be the first
                            output.push_str("\n\n");
                            // no newlines after heading needed as the message already contains them
                            output.push_str("# Assistant");
                            output.push_str(message_content.raw());
                        }
                    }
                    ConversationItemKind::AssistantReasoning(_) => {}
                    ConversationItemKind::ToolInvocation(_) => {}
                    ConversationItemKind::ToolOutput(_) => {}
                    ConversationItemKind::Notification(_, _) => {}
                }
            }
            let mut export_file = fs::File::create(&path).expect("Cannot open md file for export");
            let _ = export_file
                .write(output.as_bytes())
                .expect("Cannot write md export to file");
            let item = self.conversation.add_notification(
                NotificationType::Info,
                format!("Exported conversation to '{}'.", path.to_string_lossy()),
            );
            self.ui_tx
                .send(UiMessage::ConversationItem(item.clone()))
                .await
                .unwrap();
        }
    }

    async fn handle_quit_command(&mut self) {
        self.llm_tx.send(LlmMessage::Quit).await.unwrap();
        self.ui_tx.send(UiMessage::Quit).await.unwrap();

        // @todo do not quit if changes are unsafed and command is not force exit
        if !self.conversation.messages().is_empty() {
            // save session
            // @todo make autosafe a config option and ask otherwise
            if let Some(proj_dirs) = ProjectDirs::from("org", "codemonk", "monothink") {
                let mut state_dir = proj_dirs
                    .state_dir()
                    .expect("Running on other OS than Linux is not supported.")
                    .to_path_buf();
                state_dir.push("sessions");
                state_dir.push(format!("{}.json", self.conversation.id()));
                let session_path = state_dir.as_path();
                // @todo stream to a writer
                let output = serde_json::to_string_pretty(&self.conversation)
                    .expect("Should be able to serealize a session.");
                let mut session_file =
                    fs::File::create(session_path).expect("Cannot open session file for writing");
                let _ = session_file
                    .write(output.as_bytes())
                    .expect("Cannot write session to file");
            } else {
                // @todo probably should not panic here, as we are quitting anyhow.
                panic!("Could not get state directory for saving session.");
            };
        }

        self.quit = true;
    }

    async fn handle_view_mode_command(&mut self, view_mode: ViewMode) {
        self.ui_tx
            .send(UiMessage::ConversationViewMode(view_mode))
            .await
            .unwrap();
    }

    fn update_token_usage(&mut self, usage: Usage) {
        // assume that OpenAI typical token usage reports,
        // i.e. only input and output are reported
        if usage.input_tokens.unwrap_or_default() > 0 {
            self.token_usage.input =
                Some(self.token_usage.input.unwrap_or(0) + usage.input_tokens.unwrap_or_default());
        }
        if usage.output_tokens.unwrap_or_default() > 0 {
            self.token_usage.output = Some(
                self.token_usage.output.unwrap_or(0) + usage.output_tokens.unwrap_or_default(),
            );
        }
    }

    async fn select_model(&mut self, model: String) {
        let (model, endpoint) = if self.config.endpoint.len() > 1 {
            match model.rsplit_once("@") {
                Some((model, endpoint)) => (model.to_string(), Some(endpoint.to_string())),
                None => (model, None),
            }
        } else {
            (model, None)
        };
        // @todo check if model is valid for endpoint
        // @todo check if endpoint has changed

        self.active_model = Some(model.clone());
        self.context_window = context_window(&self.config, &model);

        self.set_model(&model).await;
    }

    async fn set_model(&self, active_model: &str) {
        self.llm_tx
            .send(LlmMessage::SetContextWindow(self.context_window))
            .await
            .unwrap();
        let provider_model_name = self
            .config
            .model
            .get(active_model)
            .and_then(|model| model.provider_model.clone())
            .unwrap_or(active_model.to_string());
        // @todo wait for confirmation from Llm
        self.ui_tx
            .send(UiMessage::ActiveModel(Some(active_model.to_string())))
            .await
            .unwrap();

        self.llm_tx
            .send(LlmMessage::SelectModel(provider_model_name))
            .await
            .unwrap();
    }

    async fn update_model_list(&self) {
        let mut suggestions = Vec::new();
        let multiple_endpoints = self.config.endpoint.len() > 1;
        for endpoint in &self.config.endpoint {
            suggestions.push(Suggestion {
                raw: if multiple_endpoints {
                    format!("{}@{}", endpoint.1.default_model, endpoint.0)
                } else {
                    endpoint.1.default_model.clone()
                },
                name: vec![Inline::Text(endpoint.1.default_model.clone())],
                hint: vec![Inline::Text(endpoint.0.to_string())],
            });
            for model in &endpoint.1.alternative_models {
                suggestions.push(Suggestion {
                    raw: if multiple_endpoints {
                        format!("{}@{}", model, endpoint.0)
                    } else {
                        model.clone()
                    },
                    name: vec![Inline::Text(model.clone())],
                    hint: vec![Inline::Text(endpoint.0.to_string())],
                });
            }
        }
        self.ui_tx
            .send(UiMessage::ModelList(suggestions))
            .await
            .unwrap();
    }
}

// @todo should be part of config
fn context_window(config: &Config, model: &str) -> u64 {
    config
        .model
        .get(model)
        .map(|model| model.context_window)
        .unwrap_or_else(default_context_window)
}