1use base64::Engine;
17use base64::engine::general_purpose::STANDARD;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum ImageFormat {
26 Png,
27 Jpeg,
28}
29
30impl ImageFormat {
31 fn kitty_format_code(self) -> &'static str {
32 match self {
33 ImageFormat::Png | ImageFormat::Jpeg => "100",
36 }
37 }
38}
39
40#[must_use]
43pub fn sniff(bytes: &[u8]) -> Option<ImageFormat> {
44 if bytes.len() >= 8 && bytes[0..8] == [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] {
45 return Some(ImageFormat::Png);
46 }
47 if bytes.len() >= 3 && bytes[0..3] == [0xFF, 0xD8, 0xFF] {
48 return Some(ImageFormat::Jpeg);
49 }
50 None
51}
52
53#[must_use]
57pub fn supports_kitty<F>(env_lookup: F) -> bool
58where
59 F: Fn(&str) -> Option<String>,
60{
61 if env_lookup("KITTY_WINDOW_ID").is_some() {
62 return true;
63 }
64 let term = env_lookup("TERM").unwrap_or_default();
65 if term.contains("kitty") {
66 return true;
67 }
68 let term_program = env_lookup("TERM_PROGRAM").unwrap_or_default();
69 matches!(term_program.as_str(), "WezTerm" | "ghostty")
70}
71
72pub fn encode_kitty(format: ImageFormat, bytes: &[u8], out: &mut String) {
83 let encoded = STANDARD.encode(bytes);
84 const CHUNK: usize = 4096;
85 let f = format.kitty_format_code();
86
87 if encoded.len() <= CHUNK {
88 out.push_str("\x1b_Ga=T,f=");
89 out.push_str(f);
90 out.push(';');
91 out.push_str(&encoded);
92 out.push_str("\x1b\\");
93 return;
94 }
95
96 let mut i = 0;
97 let mut first = true;
98 while i < encoded.len() {
99 let end = (i + CHUNK).min(encoded.len());
100 let last = end == encoded.len();
101 out.push_str("\x1b_G");
102 if first {
103 out.push_str("a=T,f=");
104 out.push_str(f);
105 out.push(',');
106 first = false;
107 }
108 out.push_str(if last { "m=0;" } else { "m=1;" });
109 out.push_str(&encoded[i..end]);
110 out.push_str("\x1b\\");
111 i = end;
112 }
113}
114
115#[must_use]
120pub fn try_render_inline<F>(bytes: &[u8], env_lookup: F) -> Option<String>
121where
122 F: Fn(&str) -> Option<String>,
123{
124 let format = sniff(bytes)?;
125 if !supports_kitty(env_lookup) {
126 return None;
127 }
128 let mut out = String::with_capacity(bytes.len() * 4 / 3 + 32);
129 encode_kitty(format, bytes, &mut out);
130 Some(out)
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use std::collections::HashMap;
137
138 fn env_from(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
139 let map: HashMap<String, String> = pairs
140 .iter()
141 .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
142 .collect();
143 move |k: &str| map.get(k).cloned()
144 }
145
146 #[test]
147 fn sniff_recognises_png_magic() {
148 let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0, 0];
149 assert_eq!(sniff(&png), Some(ImageFormat::Png));
150 }
151
152 #[test]
153 fn sniff_recognises_jpeg_magic() {
154 let jpeg = [0xFF, 0xD8, 0xFF, 0xE0, 0, 0];
155 assert_eq!(sniff(&jpeg), Some(ImageFormat::Jpeg));
156 }
157
158 #[test]
159 fn sniff_returns_none_for_random_bytes() {
160 assert_eq!(sniff(&[1, 2, 3, 4]), None);
161 assert_eq!(sniff(b"<svg ..."), None);
162 assert_eq!(sniff(&[]), None);
163 }
164
165 #[test]
166 fn supports_kitty_window_id() {
167 let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
168 assert!(supports_kitty(env));
169 }
170
171 #[test]
172 fn supports_kitty_via_term() {
173 let env = env_from(&[("TERM", "xterm-kitty")]);
174 assert!(supports_kitty(env));
175 }
176
177 #[test]
178 fn supports_kitty_via_term_program_wezterm() {
179 let env = env_from(&[("TERM_PROGRAM", "WezTerm")]);
180 assert!(supports_kitty(env));
181 }
182
183 #[test]
184 fn supports_kitty_via_term_program_ghostty() {
185 let env = env_from(&[("TERM_PROGRAM", "ghostty")]);
186 assert!(supports_kitty(env));
187 }
188
189 #[test]
190 fn supports_kitty_plain_xterm_is_false() {
191 let env = env_from(&[("TERM", "xterm-256color")]);
192 assert!(!supports_kitty(env));
193 }
194
195 #[test]
196 fn supports_kitty_empty_env_is_false() {
197 let env = env_from(&[]);
198 assert!(!supports_kitty(env));
199 }
200
201 #[test]
202 fn encode_kitty_short_payload_uses_single_chunk() {
203 let mut out = String::new();
204 encode_kitty(ImageFormat::Png, &[0x89, 0x50, 0x4E, 0x47], &mut out);
205 assert!(out.starts_with("\x1b_Ga=T,f=100;"));
207 assert!(out.ends_with("\x1b\\"));
208 assert_eq!(out.matches("\x1b_G").count(), 1);
210 }
211
212 #[test]
213 fn encode_kitty_long_payload_chunks() {
214 let bytes = vec![0xABu8; 10_000];
215 let mut out = String::new();
216 encode_kitty(ImageFormat::Png, &bytes, &mut out);
217 let starts = out.matches("\x1b_G").count();
220 assert!(starts > 1, "got {starts} chunks");
221 assert!(
222 out.ends_with(
223 "m=0;"
224 .chars()
225 .chain("".chars())
226 .collect::<String>()
227 .as_str()
228 ) || out.contains("m=0;")
229 );
230 }
231
232 #[test]
233 fn try_render_inline_returns_none_for_non_image() {
234 let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
235 assert!(try_render_inline(b"not an image", env).is_none());
236 }
237
238 #[test]
239 fn try_render_inline_returns_none_when_terminal_lacks_kitty() {
240 let env = env_from(&[("TERM", "xterm-256color")]);
241 let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
242 assert!(try_render_inline(&png, env).is_none());
243 }
244
245 #[test]
246 fn try_render_inline_returns_some_when_png_and_kitty() {
247 let env = env_from(&[("KITTY_WINDOW_ID", "1")]);
248 let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
249 let rendered = try_render_inline(&png, env).expect("should encode");
250 assert!(rendered.contains("\x1b_G"));
251 }
252}