spaceshell/crates/spaceshd/src/snapshot_store.rs

use std::path::PathBuf;
use spacesh_core::snapshot::Snapshot;
use spacesh_proto::SurfaceId;

/// Stores one visible-screen snapshot per surface as `<dir>/<surface_id>.json`.
pub trait SnapshotStore: Send + Sync {
    fn save(&self, sid: &SurfaceId, snap: &Snapshot);
    fn load(&self, sid: &SurfaceId) -> Option<Snapshot>;
    fn remove(&self, sid: &SurfaceId);
}

/// Writer command: persist or delete a surface's snapshot. Shared by the
/// router ticker, the close/remove paths, and each actor's on-exit dump, so a
/// single channel type flows everywhere.
pub enum SnapshotMsg {
    Save(SurfaceId, Snapshot),
    Remove(SurfaceId),
}

/// A no-op store for tests and contexts that do not persist snapshots.
pub struct NullSnapshotStore;
impl SnapshotStore for NullSnapshotStore {
    fn save(&self, _sid: &SurfaceId, _snap: &Snapshot) {}
    fn load(&self, _sid: &SurfaceId) -> Option<Snapshot> { None }
    fn remove(&self, _sid: &SurfaceId) {}
}

/// JSON file store. Filenames are the surface id (e.g. `s_1f.json`); ids are
/// `^[a-z]_[0-9a-f]+$` so they are always safe path components.
pub struct JsonSnapshotStore {
    dir: PathBuf,
}

impl JsonSnapshotStore {
    pub fn new(dir: PathBuf) -> Self {
        let _ = std::fs::create_dir_all(&dir);
        Self { dir }
    }
    fn path(&self, sid: &SurfaceId) -> PathBuf {
        self.dir.join(format!("{}.json", sid.0))
    }
}

impl SnapshotStore for JsonSnapshotStore {
    fn save(&self, sid: &SurfaceId, snap: &Snapshot) {
        let path = self.path(sid);
        let tmp = path.with_extension("json.tmp");
        let Ok(bytes) = serde_json::to_vec(snap) else { return };
        if std::fs::write(&tmp, &bytes).is_err() { return; }
        if std::fs::File::open(&tmp).and_then(|f| f.sync_all()).is_err() { return; }
        let _ = std::fs::rename(&tmp, &path);
    }
    fn load(&self, sid: &SurfaceId) -> Option<Snapshot> {
        let bytes = std::fs::read(self.path(sid)).ok()?;
        serde_json::from_slice(&bytes).ok()
    }
    fn remove(&self, sid: &SurfaceId) {
        let _ = std::fs::remove_file(self.path(sid));
    }
}

/// Spawn the writer task; returns the sender used by the router and actors.
pub fn spawn_writer(store: std::sync::Arc<dyn SnapshotStore>) -> tokio::sync::mpsc::UnboundedSender<SnapshotMsg> {
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<SnapshotMsg>();
    tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            match msg {
                SnapshotMsg::Save(sid, snap) => store.save(&sid, &snap),
                SnapshotMsg::Remove(sid) => store.remove(&sid),
            }
        }
    });
    tx
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tmp_dir(name: &str) -> PathBuf {
        let n = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
        let p = std::env::temp_dir().join(format!("spacesh-snap-{name}-{n}"));
        std::fs::create_dir_all(&p).unwrap();
        p
    }

    fn sample() -> Snapshot {
        Snapshot { ansi: "\u{1b}[mhello".into(), cols: 80, rows: 24, cursor_row: 1, cursor_col: 6 }
    }

    #[test]
    fn save_then_load_round_trips() {
        let dir = tmp_dir("roundtrip");
        let store = JsonSnapshotStore::new(dir.clone());
        let sid = SurfaceId("s_1".into());
        store.save(&sid, &sample());
        assert_eq!(store.load(&sid), Some(sample()));
        let _ = std::fs::remove_dir_all(dir);
    }

    #[test]
    fn missing_loads_none() {
        let store = JsonSnapshotStore::new(tmp_dir("missing"));
        assert_eq!(store.load(&SurfaceId("s_none".into())), None);
    }

    #[test]
    fn corrupt_loads_none() {
        let dir = tmp_dir("corrupt");
        let store = JsonSnapshotStore::new(dir.clone());
        let sid = SurfaceId("s_2".into());
        std::fs::write(dir.join("s_2.json"), b"{ not json").unwrap();
        assert_eq!(store.load(&sid), None);
        let _ = std::fs::remove_dir_all(dir);
    }

    #[test]
    fn remove_deletes_file() {
        let dir = tmp_dir("remove");
        let store = JsonSnapshotStore::new(dir.clone());
        let sid = SurfaceId("s_3".into());
        store.save(&sid, &sample());
        assert!(store.load(&sid).is_some());
        store.remove(&sid);
        assert_eq!(store.load(&sid), None);
        let _ = std::fs::remove_dir_all(dir);
    }

    #[test]
    fn null_store_is_inert() {
        let store = NullSnapshotStore;
        let sid = SurfaceId("s_4".into());
        store.save(&sid, &sample());
        assert_eq!(store.load(&sid), None);
        store.remove(&sid);
    }

    #[tokio::test]
    async fn writer_saves_and_removes() {
        let dir = tmp_dir("writer");
        let store: std::sync::Arc<dyn SnapshotStore> = std::sync::Arc::new(JsonSnapshotStore::new(dir.clone()));
        let tx = spawn_writer(store.clone());
        let sid = SurfaceId("s_w".into());

        tx.send(SnapshotMsg::Save(sid.clone(), sample())).unwrap();
        let mut saved = None;
        for _ in 0..50 {
            if let Some(s) = store.load(&sid) { saved = Some(s); break; }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
        assert_eq!(saved, Some(sample()));

        tx.send(SnapshotMsg::Remove(sid.clone())).unwrap();
        let mut gone = false;
        for _ in 0..50 {
            if store.load(&sid).is_none() { gone = true; break; }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
        assert!(gone, "writer should have removed the snapshot file");
        let _ = std::fs::remove_dir_all(dir);
    }
}