slug-code/src/session/mod.rs

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Write};
use std::path::PathBuf;
use uuid::Uuid;

use crate::provider::ChatMessage;

/// Metadata for a saved session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMeta {
    pub session_id: String,
    pub created_at: String,
    pub last_used_at: String,
    pub working_directory: String,
    pub model: String,
    pub summary: String,
}

/// An in-memory session handle returned on create or fork.
#[derive(Debug, Clone)]
pub struct Session {
    pub session_id: String,
    pub messages: Vec<ChatMessage>,
}

/// Manages session persistence on disk.
pub struct SessionManager {
    sessions_dir: PathBuf,
}

impl SessionManager {
    /// Create a new SessionManager, ensuring the sessions directory exists.
    pub fn new() -> Self {
        let sessions_dir = Self::sessions_dir();
        if let Err(e) = fs::create_dir_all(&sessions_dir) {
            eprintln!("Warning: could not create sessions directory: {e}");
        }
        Self { sessions_dir }
    }

    fn sessions_dir() -> PathBuf {
        dirs::home_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join(".slug")
            .join("sessions")
    }

    fn jsonl_path(&self, session_id: &str) -> PathBuf {
        self.sessions_dir.join(format!("{session_id}.jsonl"))
    }

    fn meta_path(&self, session_id: &str) -> PathBuf {
        self.sessions_dir
            .join(format!("{session_id}.meta.json"))
    }

    fn now_iso8601() -> String {
        // Use std::time to get a basic timestamp without pulling in chrono.
        // Format: seconds since UNIX epoch as a UTC ISO-8601 string.
        use std::time::{SystemTime, UNIX_EPOCH};
        let secs = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        // Produce a minimal ISO-8601-like UTC string: YYYY-MM-DDTHH:MM:SSZ
        let s = secs;
        let sec = s % 60;
        let min = (s / 60) % 60;
        let hour = (s / 3600) % 24;
        let days = s / 86400; // days since 1970-01-01
        // Convert days → calendar date (proleptic Gregorian)
        let (year, month, day) = days_to_ymd(days);
        format!(
            "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
            year, month, day, hour, min, sec
        )
    }

    /// Create a brand-new session with a fresh UUID. Does not write to disk yet.
    pub fn create_session(&self) -> Session {
        Session {
            session_id: Uuid::new_v4().to_string(),
            messages: Vec::new(),
        }
    }

    /// Write (or overwrite) the metadata file for a session.
    pub fn write_meta(&self, meta: &SessionMeta) -> Result<()> {
        let path = self.meta_path(&meta.session_id);
        let json = serde_json::to_string_pretty(meta)
            .context("failed to serialize session metadata")?;
        fs::write(&path, json).with_context(|| format!("failed to write meta file {path:?}"))?;
        Ok(())
    }

    /// Initialize metadata for a newly created session.
    pub fn init_meta(
        &self,
        session: &Session,
        working_directory: &str,
        model: &str,
        summary: &str,
    ) -> Result<()> {
        let now = Self::now_iso8601();
        let meta = SessionMeta {
            session_id: session.session_id.clone(),
            created_at: now.clone(),
            last_used_at: now,
            working_directory: working_directory.to_string(),
            model: model.to_string(),
            summary: summary.to_string(),
        };
        self.write_meta(&meta)
    }

    /// Append a single message to the JSONL file and update last_used_at.
    pub fn save_message(&self, session_id: &str, message: &ChatMessage) -> Result<()> {
        let path = self.jsonl_path(session_id);
        let mut file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)
            .with_context(|| format!("failed to open session file {path:?}"))?;
        let line = serde_json::to_string(message).context("failed to serialize message")?;
        writeln!(file, "{line}").with_context(|| format!("failed to write to {path:?}"))?;

        // Update last_used_at in meta if it exists.
        let meta_path = self.meta_path(session_id);
        if meta_path.exists() {
            if let Ok(raw) = fs::read_to_string(&meta_path) {
                if let Ok(mut meta) = serde_json::from_str::<SessionMeta>(&raw) {
                    meta.last_used_at = Self::now_iso8601();
                    let _ = self.write_meta(&meta);
                }
            }
        }

        Ok(())
    }

    /// Read all messages from a session's JSONL file.
    pub fn load_session(&self, session_id: &str) -> Result<Vec<ChatMessage>> {
        let path = self.jsonl_path(session_id);
        if !path.exists() {
            return Ok(Vec::new());
        }
        let file =
            File::open(&path).with_context(|| format!("failed to open session file {path:?}"))?;
        let reader = BufReader::new(file);
        let mut messages = Vec::new();
        for (line_no, line) in reader.lines().enumerate() {
            let line = line.with_context(|| format!("error reading line {line_no} of {path:?}"))?;
            if line.trim().is_empty() {
                continue;
            }
            let msg: ChatMessage = serde_json::from_str(&line)
                .with_context(|| format!("failed to parse line {line_no} of {path:?}"))?;
            messages.push(msg);
        }
        Ok(messages)
    }

    /// List all sessions, sorted by last_used_at descending (most recent first).
    pub fn list_sessions(&self) -> Vec<SessionMeta> {
        let mut metas = Vec::new();
        let entries = match fs::read_dir(&self.sessions_dir) {
            Ok(e) => e,
            Err(_) => return metas,
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) == Some("json")
                && path
                    .file_name()
                    .and_then(|n| n.to_str())
                    .map(|n| n.ends_with(".meta.json"))
                    .unwrap_or(false)
            {
                if let Ok(raw) = fs::read_to_string(&path) {
                    if let Ok(meta) = serde_json::from_str::<SessionMeta>(&raw) {
                        metas.push(meta);
                    }
                }
            }
        }
        // Sort most-recent first.
        metas.sort_by(|a, b| b.last_used_at.cmp(&a.last_used_at));
        metas
    }

    /// Return the most recently used session, if any.
    pub fn get_latest_session(&self) -> Option<SessionMeta> {
        self.list_sessions().into_iter().next()
    }

    /// Create a new session that starts with all messages copied from source_id.
    pub fn fork_session(&self, source_id: &str) -> Result<Session> {
        let messages = self.load_session(source_id)
            .with_context(|| format!("failed to load source session {source_id}"))?;
        let new_id = Uuid::new_v4().to_string();
        // Write all messages into the new JSONL immediately.
        for msg in &messages {
            self.save_message(&new_id, msg)?;
        }
        Ok(Session {
            session_id: new_id,
            messages,
        })
    }
}

// ---------------------------------------------------------------------------
// Minimal days-since-epoch → (year, month, day) conversion
// (avoids a chrono dependency for the timestamp helper)
// ---------------------------------------------------------------------------

fn days_to_ymd(days: u64) -> (u64, u8, u8) {
    // Algorithm: civil date from days since 1970-01-01
    // Based on Howard Hinnant's public-domain algorithm.
    let z = days as i64 + 719468;
    let era = if z >= 0 { z } else { z - 146096 } / 146097;
    let doe = z - era * 146097;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    (y as u64, m as u8, d as u8)
}