1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364

//! Pictures in a terminal, over the Kitty graphics protocol.
//!
//! A cell grid has no pixels, so an `:image` is not painted like everything
//! else here: the painter reserves the cells and records where they are, and
//! this turns that into the escape sequences that put a picture there. The
//! terminal draws it over the blank cells the grid left behind.
//!
//! PNG only, and deliberately: the file is handed over as it is, with `f=100`,
//! so nothing here decodes an image. That is the same bargain the window
//! backend makes — `libvidya` decodes PNG and nothing else, and frq only ever
//! caches PNG — and it is what keeps a terminal backend free of an image
//! library.
//!
//! Not every terminal has the protocol, and one that has not would print the
//! escape as text: a screenful of base64 where a conversation was. So this is
//! off unless the terminal is one that is known to answer — `$TERM`,
//! `$TERM_PROGRAM` and kitty's own variable say so — and `JOLT_TUI_GRAPHICS`
//! forces it either way for a terminal this does not know about yet.

use std::collections::HashMap;
use std::io::{self, Write};

use crate::screen::Rect;

/// Where one picture goes, as the painter left it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Placement {
    /// The node it belongs to, which is what makes it the same picture across
    /// frames while it scrolls.
    pub node: u32,
    /// The file, which is what makes it the same *picture* — two nodes showing
    /// one file are transmitted once.
    pub path: String,
    /// The cells it was given.
    pub area: Rect,
    /// Rows of the picture cut off the top and the bottom by whatever it is
    /// scrolling inside. A partly-visible picture is placed partly, rather
    /// than whole and over its neighbours, or not at all and flickering.
    pub crop_top: u16,
    pub crop_bottom: u16,
}

/// A picture's size in pixels, read out of the PNG header.
///
/// The eight-byte signature, then the first chunk, which the format says is
/// `IHDR`: length, type, width, height. Sixteen bytes of a file, so this is a
/// read of the head rather than a decode of the whole.
pub fn png_size(bytes: &[u8]) -> Option<(u32, u32)> {
    const SIGNATURE: [u8; 8] = [0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
    if bytes.len() < 24 || bytes[..8] != SIGNATURE || &bytes[12..16] != b"IHDR" {
        return None;
    }
    let read =
        |at: usize| u32::from_be_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]]);
    let (w, h) = (read(16), read(20));
    (w > 0 && h > 0).then_some((w, h))
}

/// Whether this terminal draws pictures.
///
/// A name rather than a question asked over the wire: the protocol's query is
/// an escape whose answer comes back as input, and the answer arrives while
/// the first frame is being painted — too late for the layout that has to
/// decide how many rows a picture takes. Names are what a terminal is
/// identified by everywhere else here, and `JOLT_TUI_GRAPHICS=1` is the way in
/// for one this list has not learned yet.
pub fn supported() -> bool {
    #[cfg(test)]
    if let Some(forced) = FORCED.with(|f| f.get()) {
        return forced;
    }
    match std::env::var("JOLT_TUI_GRAPHICS").as_deref() {
        Ok("1") | Ok("true") => return true,
        Ok("0") | Ok("false") => return false,
        _ => {}
    }
    if std::env::var_os("KITTY_WINDOW_ID").is_some()
        || std::env::var_os("GHOSTTY_RESOURCES_DIR").is_some()
    {
        return true;
    }
    let known = |name: String| {
        let name = name.to_ascii_lowercase();
        name.contains("kitty") || name.contains("ghostty") || name.contains("wezterm")
    };
    std::env::var("TERM").map(known).unwrap_or(false)
        || std::env::var("TERM_PROGRAM").map(known).unwrap_or(false)
}

/// One cell in pixels. The terminal is asked; where it will not say — a
/// multiplexer in the way, a terminal that answers zero — the usual size of a
/// cell in a terminal font stands in, which is wrong by a little rather than
/// by an order of magnitude.
pub fn cell_pixels(columns: u16, rows: u16, width: u16, height: u16) -> (u16, u16) {
    let w = if columns > 0 && width > 0 {
        width / columns
    } else {
        0
    };
    let h = if rows > 0 && height > 0 {
        height / rows
    } else {
        0
    };
    (if w > 0 { w } else { 8 }, if h > 0 { h } else { 16 })
}

thread_local! {
    /// An answer for [`supported`] that does not come from the environment —
    /// what a test says, so one asserting on a screen with a picture in it
    /// does not depend on the terminal it happens to be run from.
    #[cfg(test)]
    static FORCED: std::cell::Cell<Option<bool>> = const { std::cell::Cell::new(None) };
    /// What a cell measures, and the size of every picture on screen follows
    /// from it. The terminal is asked once when it opens and again when it is
    /// resized; a headless session keeps the default, which is what makes a
    /// test of a layout with a picture in it repeatable.
    static CELL: std::cell::Cell<(u16, u16)> = const { std::cell::Cell::new((8, 16)) };
    /// Sizes read out of PNG headers, by path. A layout asks for every picture
    /// on screen every frame, and a file's header does not change under it —
    /// frq writes a cached picture once, under a name of its own.
    static SIZES: std::cell::RefCell<HashMap<String, Option<(u32, u32)>>> =
        std::cell::RefCell::new(HashMap::new());
}

/// Say yes or no for this thread, whatever the terminal is. `None` gives the
/// question back to the environment.
#[cfg(test)]
pub fn force(answer: Option<bool>) {
    FORCED.with(|f| f.set(answer));
}

pub fn set_cell(cell: (u16, u16)) {
    CELL.with(|c| c.set(cell));
}

pub fn cell() -> (u16, u16) {
    CELL.with(|c| c.get())
}

/// The size of the picture at `path`, remembered.
pub fn size_of(path: &str) -> Option<(u32, u32)> {
    SIZES.with(|sizes| {
        if let Some(known) = sizes.borrow().get(path) {
            return *known;
        }
        let size = std::fs::read(path).ok().and_then(|bytes| png_size(&bytes));
        sizes.borrow_mut().insert(path.to_owned(), size);
        size
    })
}

/// How many cells a picture takes, given what it may not exceed.
///
/// The caller's `:max-width` and `:max-height` are the bounds — frq measures
/// those against the window, and against a terminal they arrive as cells — and
/// the picture is fitted inside them with its shape kept. A picture whose size
/// cannot be read at all is given a modest box: something is there, and the
/// line above it says what.
pub fn cells_for(path: &str, max_cols: u16, max_rows: u16) -> (u16, u16) {
    let (max_cols, max_rows) = (max_cols.max(1), max_rows.max(1));
    let Some((pw, ph)) = size_of(path) else {
        return (max_cols.min(24), max_rows.min(3));
    };
    let (cw, ch) = cell();
    let natural_cols = (pw as f64 / cw as f64).ceil().max(1.0);
    let natural_rows = (ph as f64 / ch as f64).ceil().max(1.0);
    // Whichever bound it meets first, and never scaled up: a small picture is
    // small, the way it is in the window.
    let scale = (max_cols as f64 / natural_cols)
        .min(max_rows as f64 / natural_rows)
        .min(1.0);
    (
        ((natural_cols * scale).round() as u16).max(1),
        ((natural_rows * scale).round() as u16).max(1),
    )
}

/// The pictures on screen, and what the terminal has been told about them.
#[derive(Default)]
pub struct Graphics {
    /// File -> the id it was transmitted under. A picture crosses the wire
    /// once, however many frames it is on screen for and wherever it scrolls.
    sent: HashMap<String, u32>,
    /// Node -> what it was last placed as. A placement that has not moved is
    /// left alone; the wire is quiet while a reader reads.
    placed: HashMap<u32, Placement>,
    next_id: u32,
}

impl Graphics {
    /// Bring the terminal's idea of what is on screen into line with `now`.
    ///
    /// Deletions first: a picture that has moved is deleted and placed again,
    /// and doing it in that order means the cells it used to be over are the
    /// terminal's own to redraw rather than a hole under the new placement.
    pub fn sync(
        &mut self,
        out: &mut impl Write,
        now: &[Placement],
        cell: (u16, u16),
    ) -> io::Result<()> {
        for (node, was) in std::mem::take(&mut self.placed) {
            let still = now.iter().any(|p| p.node == node && *p == was);
            if still {
                self.placed.insert(node, was);
            } else {
                write!(
                    out,
                    "\x1b_Ga=d,d=i,i={},p={},q=2\x1b\\",
                    self.id(&was.path),
                    node
                )?;
            }
        }
        for placement in now {
            if self.placed.get(&placement.node) == Some(placement) {
                continue;
            }
            self.place(out, placement, cell)?;
            self.placed.insert(placement.node, placement.clone());
        }
        Ok(())
    }

    /// Every picture forgotten, and the terminal told to drop them all. What
    /// closing a session, or leaving the alternate screen, has to do: a
    /// placement outlives the frame it was made in.
    pub fn clear(&mut self, out: &mut impl Write) -> io::Result<()> {
        self.sent.clear();
        self.placed.clear();
        write!(out, "\x1b_Ga=d,d=A,q=2\x1b\\")
    }

    fn id(&self, path: &str) -> u32 {
        self.sent.get(path).copied().unwrap_or(0)
    }

    fn place(
        &mut self,
        out: &mut impl Write,
        placement: &Placement,
        cell: (u16, u16),
    ) -> io::Result<()> {
        let Some(id) = self.transmit(out, &placement.path)? else {
            return Ok(());
        };
        let area = placement.area;
        if area.w == 0 || area.h == 0 {
            return Ok(());
        }
        // The cursor is where a placement lands, so it goes there first — and
        // `C=1` leaves it there rather than letting the picture move it, which
        // would put the next thing this writes somewhere else entirely.
        write!(out, "\x1b[{};{}H", area.y + 1, area.x + 1)?;
        let mut keys = format!(
            "a=p,i={id},p={},c={},r={},C=1,q=2",
            placement.node, area.w, area.h
        );
        if placement.crop_top > 0 || placement.crop_bottom > 0 {
            // In pixels of the source, which is what the protocol crops in.
            let top = placement.crop_top as u32 * cell.1 as u32;
            let rows = area.h as u32 * cell.1 as u32;
            keys.push_str(&format!(",y={top},h={rows}"));
        }
        write!(out, "\x1b_G{keys}\x1b\\")
    }

    /// Hand the file over, once. Answers the id it went under, or `None` for a
    /// file that is not there any more or is not a PNG — a fetch that failed
    /// leaves a link in the message, which is the right thing to be left with.
    fn transmit(&mut self, out: &mut impl Write, path: &str) -> io::Result<Option<u32>> {
        if let Some(id) = self.sent.get(path) {
            return Ok(Some(*id));
        }
        let Ok(bytes) = std::fs::read(path) else {
            return Ok(None);
        };
        if png_size(&bytes).is_none() {
            return Ok(None);
        }
        self.next_id += 1;
        let id = self.next_id;
        // Base64 in chunks the protocol's own size, each saying whether more
        // is coming. The first carries the keys; the rest carry only `m`.
        let encoded = base64(&bytes);
        let mut chunks = encoded.as_bytes().chunks(4096).peekable();
        let mut first = true;
        while let Some(chunk) = chunks.next() {
            let more = u8::from(chunks.peek().is_some());
            if first {
                write!(out, "\x1b_Ga=t,i={id},f=100,t=d,m={more},q=2;")?;
                first = false;
            } else {
                write!(out, "\x1b_Gm={more},q=2;")?;
            }
            out.write_all(chunk)?;
            write!(out, "\x1b\\")?;
        }
        self.sent.insert(path.to_owned(), id);
        Ok(Some(id))
    }
}

/// Standard base64, which is what the protocol's payload is written in.
fn base64(bytes: &[u8]) -> String {
    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
    for group in bytes.chunks(3) {
        let b = [
            group[0],
            *group.get(1).unwrap_or(&0),
            *group.get(2).unwrap_or(&0),
        ];
        let n = u32::from_be_bytes([0, b[0], b[1], b[2]]);
        let mut quad = [0u8; 4];
        for (i, slot) in quad.iter_mut().enumerate() {
            *slot = ALPHABET[(n >> (18 - 6 * i) & 0x3F) as usize];
        }
        for (i, ch) in quad.iter().enumerate() {
            // A group short of three bytes pads: two characters carry one
            // byte, three carry two.
            out.push(if i > group.len() { '=' } else { *ch as char });
        }
    }
    out
}

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

    #[test]
    fn base64_pads_the_way_the_alphabet_says() {
        assert_eq!(base64(b""), "");
        assert_eq!(base64(b"f"), "Zg==");
        assert_eq!(base64(b"fo"), "Zm8=");
        assert_eq!(base64(b"foo"), "Zm9v");
        assert_eq!(base64(b"foob"), "Zm9vYg==");
        assert_eq!(
            base64(b"any carnal pleasure."),
            "YW55IGNhcm5hbCBwbGVhc3VyZS4="
        );
    }

    #[test]
    fn a_png_header_gives_up_its_size_and_anything_else_gives_up_nothing() {
        let mut png = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
        png.extend_from_slice(&13u32.to_be_bytes());
        png.extend_from_slice(b"IHDR");
        png.extend_from_slice(&640u32.to_be_bytes());
        png.extend_from_slice(&480u32.to_be_bytes());
        assert_eq!(png_size(&png), Some((640, 480)));
        assert_eq!(png_size(b"not a picture at all"), None);
        assert_eq!(png_size(&png[..20]), None);
    }

    #[test]
    fn a_cell_is_the_screen_divided_by_the_grid_or_a_sensible_guess() {
        assert_eq!(cell_pixels(80, 24, 640, 384), (8, 16));
        assert_eq!(cell_pixels(80, 24, 0, 0), (8, 16));
        assert_eq!(cell_pixels(0, 0, 640, 384), (8, 16));
    }
}