spaceshell/crates/spacesh-core/src/snapshot.rs

use serde::{Deserialize, Serialize};
use alacritty_terminal::index::Point;
use alacritty_terminal::term::cell::Flags;
use alacritty_terminal::vte::ansi::Color;
use crate::grid::GridSurface;

/// Serializable snapshot returned by `attach`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Snapshot {
    /// ANSI byte dump suitable for `xterm.write()`.
    pub ansi: String,
    pub cols: u16,
    pub rows: u16,
    /// 1-based cursor position.
    pub cursor_row: u16,
    pub cursor_col: u16,
}

fn sgr_for_color(c: Color, foreground: bool) -> String {
    let base = if foreground { 38 } else { 48 };
    match c {
        Color::Named(named) => {
            // Map common named colors to SGR; default fg/bg reset for the rest.
            use alacritty_terminal::vte::ansi::NamedColor;
            let code = match named {
                NamedColor::Black => Some(if foreground { 30 } else { 40 }),
                NamedColor::Red => Some(if foreground { 31 } else { 41 }),
                NamedColor::Green => Some(if foreground { 32 } else { 42 }),
                NamedColor::Yellow => Some(if foreground { 33 } else { 43 }),
                NamedColor::Blue => Some(if foreground { 34 } else { 44 }),
                NamedColor::Magenta => Some(if foreground { 35 } else { 45 }),
                NamedColor::Cyan => Some(if foreground { 36 } else { 46 }),
                NamedColor::White => Some(if foreground { 37 } else { 47 }),
                NamedColor::BrightBlack => Some(if foreground { 90 } else { 100 }),
                NamedColor::BrightRed => Some(if foreground { 91 } else { 101 }),
                NamedColor::BrightGreen => Some(if foreground { 92 } else { 102 }),
                NamedColor::BrightYellow => Some(if foreground { 93 } else { 103 }),
                NamedColor::BrightBlue => Some(if foreground { 94 } else { 104 }),
                NamedColor::BrightMagenta => Some(if foreground { 95 } else { 105 }),
                NamedColor::BrightCyan => Some(if foreground { 96 } else { 106 }),
                NamedColor::BrightWhite => Some(if foreground { 97 } else { 107 }),
                _ => None, // Foreground/Background/Cursor etc. → use reset.
            };
            match code {
                Some(n) => format!("{n}"),
                None => format!("{}", if foreground { 39 } else { 49 }),
            }
        }
        Color::Indexed(i) => format!("{base};5;{i}"),
        Color::Spec(rgb) => format!("{base};2;{};{};{}", rgb.r, rgb.g, rgb.b),
    }
}

/// Serialize the visible grid into an ANSI dump.
pub fn snapshot_ansi(g: &GridSurface) -> Snapshot {
    let size = g.size();
    let term = g.term();
    let grid = term.grid();

    let mut out = String::new();
    out.push_str("\x1b[2J\x1b[H"); // clear + home

    let cols = size.cols;
    let lines = size.lines;

    // Track the last emitted attributes to avoid redundant SGR sequences.
    let mut last: Option<(Color, Color, Flags)> = None;

    for line in 0..lines {
        for col in 0..cols {
            let point = Point::new(alacritty_terminal::index::Line(line as i32), alacritty_terminal::index::Column(col));
            let cell = &grid[point];
            let cur = (cell.fg, cell.bg, cell.flags);
            if last != Some(cur) {
                let mut codes: Vec<String> = vec!["0".into()]; // reset, then re-apply
                if cell.flags.contains(Flags::BOLD) { codes.push("1".into()); }
                if cell.flags.contains(Flags::DIM) { codes.push("2".into()); }
                if cell.flags.contains(Flags::ITALIC) { codes.push("3".into()); }
                if cell.flags.contains(Flags::UNDERLINE) { codes.push("4".into()); }
                if cell.flags.contains(Flags::INVERSE) { codes.push("7".into()); }
                codes.push(sgr_for_color(cell.fg, true));
                codes.push(sgr_for_color(cell.bg, false));
                out.push_str(&format!("\x1b[{}m", codes.join(";")));
                last = Some(cur);
            }
            out.push(cell.c);
        }
        out.push_str("\r\n");
    }
    out.push_str("\x1b[0m"); // reset attributes at end

    let cursor = grid.cursor.point;
    let cursor_row = (cursor.line.0 as i64 + 1).clamp(1, lines as i64) as u16;
    let cursor_col = (cursor.column.0 as i64 + 1).clamp(1, cols as i64) as u16;
    out.push_str(&format!("\x1b[{cursor_row};{cursor_col}H"));

    Snapshot {
        ansi: out,
        cols: cols as u16,
        rows: lines as u16,
        cursor_row,
        cursor_col,
    }
}

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

    #[test]
    fn snapshot_contains_fed_text_and_is_deterministic() {
        let mut g = GridSurface::new(10, 3);
        g.feed(b"hi");
        let a = snapshot_ansi(&g);
        let b = snapshot_ansi(&g);
        assert_eq!(a.ansi, b.ansi, "snapshot must be deterministic");
        assert!(a.ansi.contains("hi"));
        assert!(a.ansi.starts_with("\x1b[2J\x1b[H"));
        assert_eq!(a.cols, 10);
        assert_eq!(a.rows, 3);
    }

    #[test]
    fn snapshot_round_trips_through_json() {
        let mut g = GridSurface::new(20, 4);
        g.feed(b"hello");
        let snap = snapshot_ansi(&g);
        let json = serde_json::to_string(&snap).unwrap();
        let back: Snapshot = serde_json::from_str(&json).unwrap();
        assert_eq!(back.ansi, snap.ansi);
        assert_eq!((back.cols, back.rows), (snap.cols, snap.rows));
        assert_eq!((back.cursor_row, back.cursor_col), (snap.cursor_row, snap.cursor_col));
    }

    #[test]
    fn cursor_is_one_based_after_input() {
        let mut g = GridSurface::new(10, 3);
        g.feed(b"abc");
        let s = snapshot_ansi(&g);
        // After 'abc' the cursor sits at column 4 (1-based) on row 1.
        assert_eq!(s.cursor_row, 1);
        assert_eq!(s.cursor_col, 4);
    }
}