gyandoot-v2/crates/gyandoot-memory/src/palace.rs

//! MemPalace — The core persistent memory store for GyanDoot agents.
//!
//! Directory structure:
//! ```text
//! ~/.gyandoot/palace/wings/<wing>/
//!   sessions/YYYY-MM-DD_transcript.md   (append-only daily logs)
//!   identity/
//!     L0_identity.md                    (short, < 157 tokens)
//!     L1_story.md                       (full narrative)
//!   drawer.md                           (append-only summary archive)
//! ```

use std::path::{Path, PathBuf};

use chrono::Utc;
use tracing::{debug, info, warn};

use crate::identity::IdentityStore;
use crate::transcript::TranscriptStore;

// ─── Error type ─────────────────────────────────────────────────────────

/// Errors that may occur during MemPalace operations.
#[derive(Debug)]
pub enum PalaceError {
    Io(std::io::Error),
    NotFound(String),
    Serialization(String),
}

impl std::fmt::Display for PalaceError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PalaceError::Io(e) => write!(f, "IO error: {e}"),
            PalaceError::NotFound(s) => write!(f, "Not found: {s}"),
            PalaceError::Serialization(s) => write!(f, "Serialization error: {s}"),
        }
    }
}

impl std::error::Error for PalaceError {}

impl From<std::io::Error> for PalaceError {
    fn from(e: std::io::Error) -> Self {
        PalaceError::Io(e)
    }
}

/// Convenience alias.
pub type PalaceResult<T> = Result<T, PalaceError>;

// ─── Palace paths ───────────────────────────────────────────────────────

/// Return the base palace root directory.
fn palace_root() -> PathBuf {
    let home = dirs::home_dir().expect("Cannot determine home directory");
    home.join(".gyandoot").join("palace")
}

/// Return the wing directory for a given wing name.
fn wing_dir(wing: &str) -> PathBuf {
    palace_root().join("wings").join(wing)
}

// ─── MemPalace ──────────────────────────────────────────────────────────

/// The MemPalace — persistent, file-backed agent memory organized by wing.
///
/// Each wing corresponds to an agent profile or deployment context.
/// Within a wing, conversations are logged as daily Markdown transcripts,
/// identity is stored as L0 (short, compact) and L1 (full story) files,
/// and summaries are archived in an append-only `drawer.md`.
pub struct MemPalace {
    wing: String,
    base_dir: PathBuf,
    sessions_dir: PathBuf,
    identity_dir: PathBuf,
    transcript_store: TranscriptStore,
    identity_store: IdentityStore,
}

impl MemPalace {
    /// Open (or create) the palace for a given wing.
    ///
    /// Creates directory structure if it does not exist.
    pub fn new(wing: &str) -> Self {
        let base = wing_dir(wing);
        let sessions = base.join("sessions");
        let identity = base.join("identity");

        // Ensure directory structure exists
        std::fs::create_dir_all(&sessions).unwrap_or_else(|e| {
            warn!("Failed to create sessions dir: {e}");
        });
        std::fs::create_dir_all(&identity).unwrap_or_else(|e| {
            warn!("Failed to create identity dir: {e}");
        });

        info!("MemPalace initialized for wing '{wing}' at {}", base.display());

        Self {
            wing: wing.to_string(),
            base_dir: base,
            sessions_dir: sessions,
            identity_dir: identity,
            transcript_store: TranscriptStore::new(),
            identity_store: IdentityStore::new(),
        }
    }

    /// Open the palace for a wing at a custom base path.
    pub fn with_path(wing: &str, base: PathBuf) -> Self {
        let sessions = base.join("sessions");
        let identity = base.join("identity");

        std::fs::create_dir_all(&sessions).unwrap_or_else(|e| {
            warn!("Failed to create sessions dir: {e}");
        });
        std::fs::create_dir_all(&identity).unwrap_or_else(|e| {
            warn!("Failed to create identity dir: {e}");
        });

        Self {
            wing: wing.to_string(),
            base_dir: base,
            sessions_dir: sessions,
            identity_dir: identity,
            transcript_store: TranscriptStore::new(),
            identity_store: IdentityStore::new(),
        }
    }

    // ── Generic storage ─────────────────────────────────────────────────

    /// Store a value as a JSON file under `sessions/`.
    ///
    /// The file is named `{key}.json` in the sessions directory.
    pub fn store(&self, key: &str, value: &serde_json::Value) -> PalaceResult<()> {
        let path = self.sessions_dir.join(format!("{key}.json"));
        let json = serde_json::to_string_pretty(value)
            .map_err(|e| PalaceError::Serialization(e.to_string()))?;
        std::fs::write(&path, &json)?;
        debug!("Stored memory key '{key}' to {}", path.display());
        Ok(())
    }

    /// Query memory by simple string matching against stored JSON files.
    ///
    /// Searches the `sessions/` directory for files whose content contains
    /// the query string. Returns up to `limit` matches.
    pub fn query(&self, query: &str, limit: usize) -> PalaceResult<Vec<String>> {
        let query_lower = query.to_lowercase();
        let mut results: Vec<(String, String)> = Vec::new(); // (filename, content snippet)

        let entries = std::fs::read_dir(&self.sessions_dir)?;
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().map_or(true, |e| e != "json") {
                continue;
            }
            let content = std::fs::read_to_string(&path)?;
            if content.to_lowercase().contains(&query_lower) {
                let snippet = if content.len() > 200 {
                    format!("{}…", &content[..200])
                } else {
                    content.clone()
                };
                results.push((
                    path.file_name()
                        .map(|n| n.to_string_lossy().to_string())
                        .unwrap_or_default(),
                    snippet,
                ));
            }
        }

        results.truncate(limit);
        Ok(results
            .into_iter()
            .map(|(name, snippet)| format!("[{name}]\n{snippet}"))
            .collect())
    }

    /// Export all palace data to a single file.
    ///
    /// The export bundles daily transcripts, identity files, and the drawer
    /// summary into one file at `path`.
    pub fn export_to(&self, path: &Path) -> PalaceResult<()> {
        let mut output = String::new();

        output.push_str(&format!("# MemPalace Export — Wing: {}\n\n", self.wing));

        // Identity
        output.push_str("## Identity\n\n");
        if let Ok(id) = std::fs::read_to_string(self.identity_dir.join("L0_identity.md")) {
            output.push_str(&id);
            output.push('\n');
        }
        if let Ok(story) = std::fs::read_to_string(self.identity_dir.join("L1_story.md")) {
            output.push_str(&story);
            output.push('\n');
        }

        // Drawer
        output.push_str("## Drawer (Summary Archive)\n\n");
        if let Ok(drawer) = std::fs::read_to_string(self.base_dir.join("drawer.md")) {
            output.push_str(&drawer);
            output.push('\n');
        }

        // Recent sessions
        output.push_str("## Recent Sessions\n\n");
        let mut session_files: Vec<_> = std::fs::read_dir(&self.sessions_dir)
            .into_iter()
            .flat_map(|d| d.flatten())
            .filter(|e| e.path().extension().map_or(false, |ext| ext == "md" || ext == "json"))
            .collect();
        session_files.sort_by_key(|e| e.path().metadata().map(|m| m.modified().ok()).ok().flatten());
        for f in session_files.iter().rev().take(10) {
            if let Ok(content) = std::fs::read_to_string(f.path()) {
                output.push_str(&format!("--- {} ---\n", f.file_name().to_string_lossy()));
                let truncated = if content.len() > 2000 {
                    format!("{}…\n[truncated {} chars]\n", &content[..2000], content.len())
                } else {
                    content
                };
                output.push_str(&truncated);
                output.push('\n');
            }
        }

        std::fs::write(path, &output)?;
        info!(
            "Exported palace wing '{}' to {} ({} bytes)",
            self.wing,
            path.display(),
            output.len()
        );
        Ok(())
    }

    // ── Transcript operations ───────────────────────────────────────────

    /// Append a message to the daily transcript file.
    ///
    /// Format: `YYYY-MM-DD_transcript.md` in the sessions directory.
    /// Each entry is timestamped and tagged with the role.
    pub async fn transcript_append(&self, role: &str, content: &str) -> PalaceResult<()> {
        let date_str = Utc::now().format("%Y-%m-%d").to_string();
        let path = self.sessions_dir.join(format!("{date_str}_transcript.md"));

        let timestamp = Utc::now().format("%H:%M:%S%.3f").to_string();
        let line = format!("\n## [{timestamp}] **{role}**\n{content}\n");

        std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)?
            .sync_all()?;

        // We write via a separate open to use std::io::Write
        use std::io::Write;
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)?;
        file.write_all(line.as_bytes())?;
        file.flush()?;

        debug!(
            "Appended {role} message ({len}B) to {}",
            path.display(),
            len = content.len()
        );
        Ok(())
    }

    /// Read recent transcript content.
    ///
    /// `max_chars` caps the total returned content. `tail_lines` limits to
    /// the last N lines of the most recent transcript file.
    pub async fn get_transcript(&self, max_chars: usize, tail_lines: usize) -> PalaceResult<String> {
        let date_str = Utc::now().format("%Y-%m-%d").to_string();
        let path = self.sessions_dir.join(format!("{date_str}_transcript.md"));

        if !path.exists() {
            return Ok(String::new());
        }

        let content = std::fs::read_to_string(&path)?;
        let lines: Vec<&str> = content.lines().collect();

        // Take the last `tail_lines` lines
        let start = lines.len().saturating_sub(tail_lines);
        let tail: String = lines[start..].join("\n");

        if tail.len() <= max_chars {
            Ok(tail)
        } else {
            let truncated: String = tail.chars().take(max_chars).collect();
            Ok(format!("{truncated}\n…[truncated]"))
        }
    }

    // ── Identity operations ─────────────────────────────────────────────

    /// Set the agent identity, writing both L0 (compact) and L1 (full story)
    /// files.
    ///
    /// **L0** must fit in roughly 157 tokens (~628 chars) for fast
    /// injection into every system prompt.
    pub fn identity_set(
        &self,
        name: &str,
        role: &str,
        user_name: &str,
        user_role: &str,
    ) -> PalaceResult<()> {
        // L0 identity: compact, < 157 tokens (~628 chars)
        let l0 = format!(
            "name: {name}\nrole: {role}\nuser: {user_name} ({user_role})\n"
        );

        if estimate_tokens(&l0) > 157 {
            warn!(
                "L0 identity may exceed 157 tokens (estimated {})",
                estimate_tokens(&l0)
            );
        }

        let l0_path = self.identity_dir.join("L0_identity.md");
        std::fs::write(&l0_path, &l0)?;

        // L1 story: full narrative
        let l1 = format!(
            "# {name}\n\n**Role:** {role}\n\n**User:** {user_name}\n\n**User Role:** {user_role}\n\n\
             This agent assists {user_name} ({user_role}) as {name}, playing the role of {role}.\n\n\
             ## Traits\n- Helpful\n- Knowledgeable\n- Direct\n\n## Preferences\n- Concise responses\n\
             - Action-oriented\n- Evidence-based\n"
        );
        let l1_path = self.identity_dir.join("L1_story.md");
        std::fs::write(&l1_path, &l1)?;

        info!("Identity set: name={name}, role={role}, user={user_name}");
        Ok(())
    }

    /// Read current identity (L0 + L1) as a tuple.
    pub fn identity_get(&self) -> PalaceResult<(String, String)> {
        let l0_path = self.identity_dir.join("L0_identity.md");
        let l1_path = self.identity_dir.join("L1_story.md");

        let l0 = if l0_path.exists() {
            std::fs::read_to_string(&l0_path)?
        } else {
            String::new()
        };

        let l1 = if l1_path.exists() {
            std::fs::read_to_string(&l1_path)?
        } else {
            String::new()
        };

        Ok((l0, l1))
    }

    // ── Drawer / session digest ─────────────────────────────────────────

    /// Append a summary/digest of the current session to `drawer.md`.
    ///
    /// The drawer is an append-only file that accumulates session-level
    /// summaries for long-term recall.
    pub async fn save_session_digest(&self, summary: &str) -> PalaceResult<()> {
        let path = self.base_dir.join("drawer.md");
        let timestamp = Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string();

        use std::io::Write;
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)?;
        let entry = format!("\n---\n**Session Digest — {timestamp}**\n\n{summary}\n");
        file.write_all(entry.as_bytes())?;
        file.flush()?;

        info!("Session digest saved to drawer.md");
        Ok(())
    }

    // ── Getters ─────────────────────────────────────────────────────────

    /// Get a reference to the in-memory transcript store.
    pub fn transcript_store(&self) -> &TranscriptStore {
        &self.transcript_store
    }

    /// Get a mutable reference to the in-memory transcript store.
    pub fn transcript_store_mut(&mut self) -> &mut TranscriptStore {
        &mut self.transcript_store
    }

    /// Get a reference to the identity store.
    pub fn identity_store(&self) -> &IdentityStore {
        &self.identity_store
    }

    /// Get a mutable reference to the identity store.
    pub fn identity_store_mut(&mut self) -> &mut IdentityStore {
        &mut self.identity_store
    }

    /// The wing name for this palace.
    pub fn wing(&self) -> &str {
        &self.wing
    }

    /// The base directory of this palace.
    pub fn base_dir(&self) -> &Path {
        &self.base_dir
    }
}

// ─── Helpers ────────────────────────────────────────────────────────────

/// Rough token estimation (~4 chars per token).
fn estimate_tokens(text: &str) -> usize {
    (text.len() / 4).max(1)
}

// ─── Tests ──────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    fn test_wing() -> String {
        static COUNTER: AtomicUsize = AtomicUsize::new(0);
        format!("test_wing_{}", COUNTER.fetch_add(1, Ordering::SeqCst))
    }

    #[test]
    fn test_new_creates_directories() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        assert!(palace.sessions_dir.exists());
        assert!(palace.identity_dir.exists());
        assert_eq!(palace.wing(), &wing);

        // Cleanup
        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[test]
    fn test_store_and_query() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        palace
            .store("test_key", &serde_json::json!({"content": "Python is great"}))
            .unwrap();

        let results = palace.query("Python", 10).unwrap();
        assert!(!results.is_empty());
        assert!(results[0].contains("Python"));

        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[tokio::test]
    async fn test_transcript_append_and_read() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        palace
            .transcript_append("user", "Hello there!")
            .await
            .unwrap();
        palace
            .transcript_append("assistant", "Hi! How can I help?")
            .await
            .unwrap();

        let transcript = palace.get_transcript(4096, 50).await.unwrap();
        assert!(transcript.contains("Hello there!"));
        assert!(transcript.contains("Hi! How can I help?"));

        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[test]
    fn test_identity_set_and_get() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        palace
            .identity_set("GyanDoot", "AI Assistant", "Alice", "Developer")
            .unwrap();

        let (l0, l1) = palace.identity_get().unwrap();
        assert!(l0.contains("GyanDoot"));
        assert!(l0.contains("Alice"));
        assert!(l0.contains("Developer"));
        assert!(l0.contains("AI Assistant"));
        assert!(l1.contains("GyanDoot"));

        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[tokio::test]
    async fn test_session_digest() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        palace
            .save_session_digest("Discussed project architecture and planned refactoring.")
            .await
            .unwrap();

        let drawer_path = palace.base_dir().join("drawer.md");
        assert!(drawer_path.exists());
        let content = std::fs::read_to_string(drawer_path).unwrap();
        assert!(content.contains("Session Digest"));
        assert!(content.contains("refactoring"));

        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[test]
    fn test_export_to() {
        let wing = test_wing();
        let palace = MemPalace::new(&wing);

        let tmp = std::env::temp_dir().join(format!("gyandoot_export_{wing}.md"));
        palace.export_to(&tmp).unwrap();
        assert!(tmp.exists());
        let content = std::fs::read_to_string(&tmp).unwrap();
        assert!(content.contains(&wing));
        std::fs::remove_file(&tmp).ok();
        std::fs::remove_dir_all(palace.base_dir()).ok();
    }

    #[test]
    fn test_with_path() {
        let wing = test_wing();
        let tmp = std::env::temp_dir().join(format!("gyandoot_palace_{wing}"));
        let palace = MemPalace::with_path(&wing, tmp.clone());

        assert!(palace.sessions_dir.exists());
        assert!(palace.identity_dir.exists());

        std::fs::remove_dir_all(&tmp).ok();
    }
}