rustGPT/src/core/session.rs

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;

use crate::config::Config;

const SYSTEM_PROMPT: &str = "You are an AI assistant running in a terminal (CLI) environment. \
Optimise all answers for 80‑column readability, prefer plain text, \
ASCII art or concise bullet lists over heavy markup, and wrap code \
snippets in fenced blocks when helpful. Do not emit trailing spaces or \
control characters.";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: String,
    pub content: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionData {
    pub model: String,
    pub messages: Vec<Message>,
    pub enable_web_search: bool,
    pub enable_reasoning_summary: bool,
    #[serde(default = "default_reasoning_effort")]
    pub reasoning_effort: String,
    pub updated_at: DateTime<Utc>,
}

fn default_reasoning_effort() -> String {
    "medium".to_string()
}

#[derive(Debug, Clone)]
pub struct Session {
    pub name: String,
    pub model: String,
    pub messages: Vec<Message>,
    pub enable_web_search: bool,
    pub enable_reasoning_summary: bool,
    pub reasoning_effort: String,
}

impl Session {
    pub fn new(name: String, model: String) -> Self {
        let mut session = Self {
            name,
            model,
            messages: Vec::new(),
            enable_web_search: true,
            enable_reasoning_summary: false,
            reasoning_effort: "medium".to_string(),
        };
        
        // Add system prompt as first message
        session.messages.push(Message {
            role: "system".to_string(),
            content: SYSTEM_PROMPT.to_string(),
        });
        
        session
    }

    pub fn sessions_dir() -> Result<PathBuf> {
        let config = Config::load().unwrap_or_default();
        let home = dirs::home_dir().context("Could not find home directory")?;
        let sessions_dir = home.join(&config.session.sessions_dir_name);
        
        if !sessions_dir.exists() {
            fs::create_dir_all(&sessions_dir)
                .with_context(|| format!("Failed to create sessions directory: {:?}", sessions_dir))?;
        }
        
        Ok(sessions_dir)
    }

    pub fn session_path(name: &str) -> Result<PathBuf> {
        let config = Config::load().unwrap_or_default();
        Ok(Self::sessions_dir()?.join(format!("{}.{}", name, config.session.file_extension)))
    }

    pub fn save(&self) -> Result<()> {
        let data = SessionData {
            model: self.model.clone(),
            messages: self.messages.clone(),
            enable_web_search: self.enable_web_search,
            enable_reasoning_summary: self.enable_reasoning_summary,
            reasoning_effort: self.reasoning_effort.clone(),
            updated_at: Utc::now(),
        };

        let path = Self::session_path(&self.name)?;
        let tmp_path = path.with_extension("tmp");
        
        let json_data = serde_json::to_string_pretty(&data)
            .context("Failed to serialize session data")?;
        
        fs::write(&tmp_path, json_data)
            .with_context(|| format!("Failed to write session to {:?}", tmp_path))?;
        
        fs::rename(&tmp_path, &path)
            .with_context(|| format!("Failed to rename {:?} to {:?}", tmp_path, path))?;
        
        Ok(())
    }

    pub fn load(name: &str) -> Result<Self> {
        let path = Self::session_path(name)?;
        
        if !path.exists() {
            return Err(anyhow::anyhow!("Session '{}' does not exist", name));
        }

        let json_data = fs::read_to_string(&path)
            .with_context(|| format!("Failed to read session from {:?}", path))?;
        
        let data: SessionData = serde_json::from_str(&json_data)
            .with_context(|| format!("Failed to parse session data from {:?}", path))?;

        let mut session = Self {
            name: name.to_string(),
            model: data.model,
            messages: data.messages,
            enable_web_search: data.enable_web_search,
            enable_reasoning_summary: data.enable_reasoning_summary,
            reasoning_effort: data.reasoning_effort,
        };

        // Ensure system prompt is present
        if session.messages.is_empty() || session.messages[0].role != "system" {
            session.messages.insert(0, Message {
                role: "system".to_string(),
                content: SYSTEM_PROMPT.to_string(),
            });
        }

        Ok(session)
    }

    pub fn add_user_message(&mut self, content: String) {
        self.messages.push(Message {
            role: "user".to_string(),
            content,
        });
    }

    pub fn add_assistant_message(&mut self, content: String) {
        self.messages.push(Message {
            role: "assistant".to_string(),
            content,
        });
    }

    pub fn clear_messages(&mut self) {
        self.messages.clear();
        // Re-add system prompt
        self.messages.push(Message {
            role: "system".to_string(),
            content: SYSTEM_PROMPT.to_string(),
        });
    }

    pub fn list_sessions() -> Result<Vec<(String, DateTime<Utc>)>> {
        let sessions_dir = Self::sessions_dir()?;
        
        if !sessions_dir.exists() {
            return Ok(Vec::new());
        }

        let mut sessions = Vec::new();
        
        for entry in fs::read_dir(&sessions_dir)? {
            let entry = entry?;
            let path = entry.path();
            
            if let Some(extension) = path.extension() {
                if extension == "json" {
                    if let Some(name) = path.file_stem().and_then(|s| s.to_str()) {
                        let metadata = entry.metadata()?;
                        let modified = metadata.modified()?;
                        let datetime = DateTime::<Utc>::from(modified);
                        sessions.push((name.to_string(), datetime));
                    }
                }
            }
        }
        
        sessions.sort_by(|a, b| b.1.cmp(&a.1)); // Sort by modification time, newest first
        Ok(sessions)
    }

    pub fn delete_session(name: &str) -> Result<()> {
        let path = Self::session_path(name)?;
        
        if !path.exists() {
            return Err(anyhow::anyhow!("Session '{}' does not exist", name));
        }

        fs::remove_file(&path)
            .with_context(|| format!("Failed to delete session file: {:?}", path))?;
        
        Ok(())
    }
}