use crate::error::{ViuError, ViuResult};
use crate::printer::{adjust_offset, find_best_fit, Printer};
use crate::Config;
use base64::{engine::general_purpose, Engine};
use console::{Key, Term};
use lazy_static::lazy_static;
use std::io::Write;
use std::io::{Error, ErrorKind};
pub struct KittyPrinter;
const TEMP_FILE_PREFIX: &str = ".tty-graphics-protocol.viuer.";
lazy_static! {
static ref KITTY_SUPPORT: KittySupport = check_kitty_support();
}
pub fn get_kitty_support() -> KittySupport {
*KITTY_SUPPORT
}
impl Printer for KittyPrinter {
fn print(
&self,
stdout: &mut impl Write,
img: &image::DynamicImage,
config: &Config,
) -> ViuResult<(u32, u32)> {
match get_kitty_support() {
KittySupport::None => Err(ViuError::KittyNotSupported),
KittySupport::Local => {
print_local(stdout, img, config)
}
KittySupport::Remote => {
print_remote(stdout, img, config)
}
}
}
}
#[derive(PartialEq, Eq, Copy, Clone)]
pub enum KittySupport {
None,
Local,
Remote,
}
fn check_kitty_support() -> KittySupport {
if let Ok(term) = std::env::var("TERM") {
if term.contains("kitty") {
if has_local_support().is_ok() {
return KittySupport::Local;
} else {
return KittySupport::Remote;
}
}
}
KittySupport::None
}
fn has_local_support() -> ViuResult {
let x = image::RgbaImage::new(1, 1);
let raw_img = x.as_raw();
let path = store_in_tmp_file(raw_img)?;
print!(
"\x1b_Gi=31,s=1,v=1,a=q,t=t;{}\x1b\\",
general_purpose::STANDARD.encode(path.to_str().ok_or_else(|| std::io::Error::new(
std::io::ErrorKind::Other,
"Could not convert path to &str"
))?)
);
std::io::stdout().flush()?;
let term = Term::stdout();
let mut response = Vec::new();
while let Ok(key) = term.read_key() {
let should_break = key == Key::UnknownEscSeq(vec!['\\']) || key == Key::Unknown;
response.push(key);
if should_break {
break;
}
}
let expected = [
Key::Char('O'),
Key::Char('K'),
Key::UnknownEscSeq(vec!['\\']),
];
if response.len() >= expected.len() && response[response.len() - 3..] == expected {
return Ok(());
}
Err(ViuError::KittyResponse(response))
}
fn print_local(
stdout: &mut impl Write,
img: &image::DynamicImage,
config: &Config,
) -> ViuResult<(u32, u32)> {
let rgba = img.to_rgba8();
let raw_img = rgba.as_raw();
let path = store_in_tmp_file(raw_img)?;
adjust_offset(stdout, config)?;
let (w, h) = find_best_fit(img, config.width, config.height);
write!(
stdout,
"\x1b_Gf=32,s={},v={},c={},r={},a=T,t=t;{}\x1b\\",
img.width(),
img.height(),
w,
h,
general_purpose::STANDARD.encode(path.to_str().ok_or_else(|| ViuError::Io(Error::new(
ErrorKind::Other,
"Could not convert path to &str"
)))?)
)?;
writeln!(stdout)?;
stdout.flush()?;
Ok((w, h))
}
fn print_remote(
stdout: &mut impl Write,
img: &image::DynamicImage,
config: &Config,
) -> ViuResult<(u32, u32)> {
let rgba = img.to_rgba8();
let raw = rgba.as_raw();
let encoded = general_purpose::STANDARD.encode(raw);
let mut iter = encoded.chars().peekable();
adjust_offset(stdout, config)?;
let (w, h) = find_best_fit(img, config.width, config.height);
let first_chunk: String = iter.by_ref().take(4096).collect();
write!(
stdout,
"\x1b_Gf=32,a=T,t=d,s={},v={},c={},r={},m=1;{}\x1b\\",
img.width(),
img.height(),
w,
h,
first_chunk
)?;
while iter.peek().is_some() {
let chunk: String = iter.by_ref().take(4096).collect();
let m = if iter.peek().is_some() { 1 } else { 0 };
write!(stdout, "\x1b_Gm={};{}\x1b\\", m, chunk)?;
}
writeln!(stdout)?;
stdout.flush()?;
Ok((w, h))
}
fn store_in_tmp_file(buf: &[u8]) -> std::result::Result<std::path::PathBuf, ViuError> {
let (mut tmpfile, path) = tempfile::Builder::new()
.prefix(TEMP_FILE_PREFIX)
.rand_bytes(1)
.tempfile()?
.keep()?;
tmpfile.write_all(buf)?;
tmpfile.flush()?;
Ok(path)
}
#[cfg(test)]
mod tests {
use super::*;
use image::{DynamicImage, GenericImage};
#[test]
fn test_print_local() {
let img = DynamicImage::ImageRgba8(image::RgbaImage::new(40, 25));
let config = Config {
x: 4,
y: 3,
..Default::default()
};
let mut vec = Vec::new();
assert_eq!(print_local(&mut vec, &img, &config).unwrap(), (40, 13));
let result = std::str::from_utf8(&vec).unwrap();
assert!(result.starts_with("\x1b[4;5H\x1b_Gf=32,s=40,v=25,c=40,r=13,a=T,t=t;"));
assert!(result.ends_with("\x1b\\\n"));
}
#[test]
fn test_print_remote() {
let mut img = DynamicImage::ImageRgba8(image::RgbaImage::new(1, 2));
img.put_pixel(0, 1, image::Rgba([2, 4, 6, 8]));
let config = Config {
x: 2,
y: 5,
..Default::default()
};
let mut vec = Vec::new();
assert_eq!(print_remote(&mut vec, &img, &config).unwrap(), (1, 1));
let result = std::str::from_utf8(&vec).unwrap();
assert_eq!(
result,
"\x1b[6;3H\x1b_Gf=32,a=T,t=d,s=1,v=2,c=1,r=1,m=1;AAAAAAIEBgg=\x1b\\\n"
);
}
}