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

use std::{
    env,
    ffi::OsStr,
    fs::{self, File},
    io::{Read, Write},
    path::{Path, PathBuf},
    time::{Duration, SystemTime},
};

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

use crate::{
    common::{
        conversation::{
            Conversation, ConversationItem, ConversationItemKind, Inline, MessageState,
            NotificationType,
        },
        sender::Sender,
    },
    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),
    LoadSession(String),
    Quit,
    ResendPrompt,
    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,
    },
    FinalAgentResponse(String),
    Notification(NotificationType, String),
    ToolCall {
        name: String,
        args: serde_json::Map<String, serde_json::Value>,
        call_id: CallId,
    },
    ToolResult {
        name: String,
        output: String,
        success: bool,
        call_id: CallId,
    },
}

#[derive(Debug)]
pub struct Session {
    // set after the user submitted a prompt or changed the title
    // used to decide if the session needs to be persisted
    changed_significantly: bool,
    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 {
            changed_significantly: false,
            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,
        }
    }

    #[tracing::instrument(skip(self))]
    pub async fn run(&mut self) {
        tracing::info!("Starting session actor");
        self.clean_sessions().await;
        self.update_session_list().await;
        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;
        if let Some(active_model) = self.active_model.clone() {
            // @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;
        self.ui_tx
            .send(UiMessage::ContextState {
                fill_state: 0,
                total: self.context_window,
            })
            .await;

        while !self.quit {
            if self.ui_tx.closed().await || self.llm_tx.closed().await {
                let ui_channel_state = self.ui_tx.closed().await;
                let llm_channel_state = self.llm_tx.closed().await;
                tracing::warn!(
                    ui_tx = ui_channel_state,
                    llm_tx = llm_channel_state,
                    "Shutting down session actor because of a closed channel"
                );
                self.quit = true;
                continue;
            }
            let event = self.rx.recv().await;
            if let Some(msg) = event {
                match msg {
                    SessionMessage::SubmitCommand(command) => self.evaluate_command(command).await,
                    SessionMessage::SubmitPrompt(prompt) => self.send_prompt(prompt).await,
                    SessionMessage::AgentResponse(chunk) => {
                        let item = self.conversation.add_assistant_response_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await;
                    }
                    SessionMessage::AgentReasoning(chunk) => {
                        let item = self.conversation.add_assistant_reasoning_chunk(&chunk);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await;
                    }
                    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;

                        // 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;
                    }
                    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;
                        self.ui_tx
                            .send(UiMessage::ContextState {
                                fill_state: updated,
                                total: window,
                            })
                            .await;
                    }
                    SessionMessage::FinalAgentResponse(msg) => {
                        let item = self.conversation.set_final_assistant_response(&msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await;
                        self.ui_tx.send(UiMessage::IsProcessing(false)).await;
                    }
                    SessionMessage::Notification(nt, msg) => {
                        let item = self.conversation.add_notification(nt, msg);
                        self.ui_tx
                            .send(UiMessage::ConversationItem(item.to_owned()))
                            .await;
                        self.ui_tx.send(UiMessage::IsProcessing(false)).await;
                    }
                    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;
                    }
                    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;
                    }
                }
            }
        }
    }

    #[tracing::instrument(skip(self))]
    async fn send_prompt(&mut self, prompt: String) {
        let item = self.conversation.add_user_prompt(&prompt).clone();
        self.llm_tx
            .send(LlmMessage::SubmitPrompt {
                prompt,
                history: self.conversation.messages().clone(),
            })
            .await;
        self.ui_tx
            .send(UiMessage::ConversationItem(item.clone()))
            .await;
        self.ui_tx.send(UiMessage::IsProcessing(true)).await;
        self.changed_significantly = true;
    }

    #[tracing::instrument(skip(self))]
    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::LoadSession(id) => self.handle_load_session_command(&id).await,
            Command::Quit => self.handle_quit_command().await,
            Command::ResendPrompt => self.handle_resend_prompt().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;
                self.changed_significantly = true;
            }
            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;
        }
    }

    async fn handle_load_session_command(&mut self, id: &str) {
        let Some(mut session_path) = self.get_session_directory().await else {
            // @todo log event (which is okay, because session dir might not exist)
            return;
        };
        session_path.push(format!("{}.json", id));
        self.conversation.clear();
        self.conversation.load(session_path);

        self.ui_tx.send(UiMessage::ClearConversation).await;
        for message in self.conversation.messages() {
            self.ui_tx
                .send(UiMessage::ConversationItem(message.clone()))
                .await;
        }
        if let Some(title) = self.conversation.title() {
            self.ui_tx.send(UiMessage::Title(title)).await;
        }
    }

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

        if self.changed_significantly {
            // save session
            if let Some(mut state_dir) = self.get_session_directory().await {
                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_resend_prompt(&mut self) {
        if let Some(ConversationItem {
            content: ConversationItemKind::UserMessage(message),
            ..
        }) = self
            .conversation
            .messages()
            .iter()
            .rev()
            .find(|msg| msg.is_user_message())
        {
            self.send_prompt(message.raw().clone()).await
        }
        // @todo send error
    }

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

    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 and set it

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

        self.set_model(&model).await;
    }

    async fn set_model(&mut self, active_model: &str) {
        self.llm_tx
            .send(LlmMessage::SetContextWindow(self.context_window))
            .await;
        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.llm_tx
            .send(LlmMessage::SelectModel(provider_model_name))
            .await;

        self.ui_tx
            .send(UiMessage::ActiveModel(Some(active_model.to_string())))
            .await;
        let item = self.conversation.add_notification(
            NotificationType::Info,
            format!("Set active model to '{}'", active_model),
        );
        self.ui_tx
            .send(UiMessage::ConversationItem(item.clone()))
            .await;
    }

    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;
    }

    async fn clean_sessions(&mut self) {
        let duration_untitled = match self.config.general.untitled_session_cleanup {
            crate::config::SessionCleanup::Days(d) => Duration::from_hours(d * 24), //std::time::Duration::from_days(d),
            crate::config::SessionCleanup::Month(m) => Duration::from_hours(m * 24 * 30), //Duration::from_month(m),
            crate::config::SessionCleanup::Weeks(w) => Duration::from_hours(w * 24 * 7), //Duration::from_weeks(w),
            crate::config::SessionCleanup::Immediate => Duration::from_hours(0),
            crate::config::SessionCleanup::Never => return,
        };
        let duration_titled = match self.config.general.titled_session_cleanup {
            crate::config::SessionCleanup::Days(d) => Duration::from_hours(d * 24), //std::time::Duration::from_days(d),
            crate::config::SessionCleanup::Month(m) => Duration::from_hours(m * 24 * 30), //Duration::from_month(m),
            crate::config::SessionCleanup::Weeks(w) => Duration::from_hours(w * 24 * 7), //Duration::from_weeks(w),
            crate::config::SessionCleanup::Immediate => Duration::from_hours(0),
            crate::config::SessionCleanup::Never => return,
        };

        let Some(session_dir) = self.get_session_directory().await else {
            // @todo log else case. It is okay, as the directory could not exist on first start,
            // but we should log it to know about it
            return;
        };
        let Ok(entries) = session_dir.read_dir() else {
            // @todo log error
            return;
        };

        let mut num_deleted = 0;
        let now = SystemTime::now();
        for maybe_entry in entries {
            let Ok(entry) = maybe_entry else {
                // @todo log
                continue;
            };

            if entry.path().is_file() && entry.path().extension() == Some(OsStr::new("json")) {
                let Ok(mut file) = File::open(entry.path()) else {
                    // @todo log error
                    continue;
                };

                let has_title = {
                    let mut contents = String::new();
                    let result = file.read_to_string(&mut contents);
                    if result.is_err() {
                        // @todo log error
                        continue;
                    }
                    let Ok(conversation): Result<Conversation, serde_json::Error> =
                        serde_json::from_str(&contents)
                    else {
                        // @todo log error
                        continue;
                    };
                    conversation.title().is_some()
                };
                let Ok(modified) = file.metadata().and_then(|md| md.modified()) else {
                    // @todo log error
                    continue;
                };

                let duration = if has_title {
                    duration_titled
                } else {
                    duration_untitled
                };

                if modified + duration < now {
                    let result = std::fs::remove_file(entry.path());
                    if result.is_ok() {
                        num_deleted += 1;
                    }
                }
            }
        }

        if num_deleted > 0 {
            let item = self.conversation.add_notification(
                NotificationType::Debug,
                format!("Cleared up {} old sessions.", num_deleted),
            );
            self.ui_tx
                .send(UiMessage::ConversationItem(item.clone()))
                .await;
        }
    }

    async fn update_session_list(&self) {
        let Some(session_dir) = self.get_session_directory().await else {
            // @todo log else case. It is okay, as the directory could not exist on first start,
            // but we should log it to know about it
            return;
        };
        let Ok(entries) = session_dir.read_dir() else {
            // @todo log error
            return;
        };

        let mut suggestions = Vec::new();
        for maybe_entry in entries {
            let Ok(entry) = maybe_entry else {
                // @todo log
                continue;
            };

            if entry.path().is_file() && entry.path().extension() == Some(OsStr::new("json")) {
                let Ok(mut file) = File::open(entry.path()) else {
                    // @todo log error
                    continue;
                };
                let mut contents = String::new();
                let result = file.read_to_string(&mut contents);
                if result.is_err() {
                    // @todo log error
                    continue;
                }
                let Ok(conversation): Result<Conversation, serde_json::Error> =
                    serde_json::from_str(&contents)
                else {
                    // @todo log error
                    continue;
                };

                if self.config.general.filter_sessions_by_workspace
                    && !conversation
                        .workspace_root()
                        .starts_with(self.conversation.workspace_root())
                {
                    continue;
                }
                let Ok(modified) = file.metadata().and_then(|md| md.modified()) else {
                    // @todo log error
                    continue;
                };
                let Ok(modified_text) = std::convert::Into::<OffsetDateTime>::into(modified)
                    .format(format_description!("[year]-[month]-[day] [hour]:[minute]"))
                else {
                    // @todo log error
                    continue;
                };
                if !self.config.general.show_untitled_sessions && conversation.title().is_none() {
                    continue;
                }
                let name = match conversation.title() {
                    Some(title) => vec![Inline::Text(format!("{} ({})", title, modified_text))],
                    None => vec![Inline::Text(format!("<NO TITLE> ({})", modified_text))],
                };
                suggestions.push(Suggestion {
                    raw: conversation.id().to_string(),
                    name,
                    hint: Vec::new(),
                });
            }
        }
        self.ui_tx.send(UiMessage::SessionList(suggestions)).await;
    }

    async fn get_session_directory(&self) -> Option<PathBuf> {
        if let Some(proj_dirs) = ProjectDirs::from("org", "codemonk", "monothink") {
            let mut state_dir = proj_dirs
                .state_dir()
                .unwrap_or_else(|| proj_dirs.data_dir())
                .to_path_buf();
            state_dir.push("sessions");
            Some(state_dir)
        } else {
            None
        }
    }
}

// @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)
}