tuigram-cli 0.9.0

//! Photos, shown inline in chat bubbles, and chat photos in the chat list.
//!
//! Every frame, the message view and the chat list ask for the photos they
//! have on screen. After the frame, [`Images::fetch`] starts whatever is
//! missing: a TDLib download, then decode + encode for the terminal on a
//! blocking thread. Until the real photo is ready, the blurry thumbnail that
//! came with the message (or chat) stands in.

use std::cell::Cell;
use std::collections::{HashMap, HashSet};
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::Arc;

use anyhow::Result;
use image::imageops::FilterType;
use image::{DynamicImage, RgbaImage};
use ratatui::layout::Size;
use ratatui_image::picker::{Picker, ProtocolType};
use ratatui_image::sliced::SlicedProtocol;
use ratatui_image::{FontSize, Resize};
use tokio::sync::mpsc::UnboundedSender;

use crate::chats::ChatPhoto;
use crate::messages::Preview;
use crate::tg::Tg;

/// Scale both ways (`Fit` never enlarges, which would leave the tiny thumbnail
/// tiny). Triangle filtering keeps the enlarged thumbnail soft, not blocky.
const RESIZE: Resize = Resize::Scale(Some(FilterType::Triangle));

/// Chat photos kept encoded; the ones drawn longest ago go first. Enough for
/// several screens of the list.
const MAX_AVATARS: usize = 300;

/// One encoded image: a photo (or its thumbnail) at one size in cells.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct Key {
    pub file_id: i32,
    pub cols: u16,
    pub rows: u16,
    /// The blurry thumbnail embedded in the message, not the downloaded photo.
    pub thumbnail: bool,
    /// A chat photo for the chat list, cut to a circle.
    pub avatar: bool,
}

/// Caps for decoding images from other people. Telegram's photos are at most
/// 2560px a side and stickers 512px, so a bigger one is only a way to run
/// memory out.
fn limits() -> image::Limits {
    sized(4096)
}

/// Caps of `side` pixels a side. Decoders allocate more than `max_alloc`
/// counts (WebP's scratch and frame buffers), so the side is what bounds it.
fn sized(side: u32) -> image::Limits {
    let mut limits = image::Limits::default();
    limits.max_image_width = Some(side);
    limits.max_image_height = Some(side);
    limits.max_alloc = Some(128 * 1024 * 1024);
    limits
}

/// Stickers are 512px a side; this leaves room to spare.
const STICKER_SIDE: u32 = 1024;

/// Like `image::open` (the format comes from the extension), within [`limits`].
pub fn open_image(path: &str) -> image::ImageResult<image::DynamicImage> {
    open_within(path, limits())
}

fn open_within(path: &str, limits: image::Limits) -> image::ImageResult<image::DynamicImage> {
    let mut reader = image::ImageReader::open(path)?;
    reader.limits(limits);
    reader.decode()
}

/// At most this many images decode at once; the rest wait for a later frame.
/// Each may take a few hundred MB while it decodes.
const MAX_BUILDING: usize = 4;

/// Encoded photos kept for the open chat; the least recently drawn go first.
const MAX_READY: usize = 200;

thread_local! {
    /// Set while this thread decodes another user's image.
    static DECODING: Cell<bool> = const { Cell::new(false) };
}

/// The panic is in an image decoder, where it's caught: the image shows as
/// broken and the app carries on, rather than every chat holding that image
/// becoming one that can't be opened.
pub fn panic_is_contained() -> bool {
    DECODING.with(Cell::get)
}

/// Runs `decode`, with a panic in it caught and turned into an error. Every
/// decode of another user's image goes through this.
pub fn contained<T>(decode: impl FnOnce() -> Result<T>) -> Result<T> {
    DECODING.with(|d| d.set(true));
    let result = catch_unwind(AssertUnwindSafe(decode))
        .unwrap_or_else(|_| Err(anyhow::anyhow!("the image couldn't be decoded")));
    DECODING.with(|d| d.set(false));
    result
}

/// An image still building after this long counts as failed.
const BUILD_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);

/// Like `image::load_from_memory`, within [`limits`].
fn decode_bytes(data: &[u8]) -> image::ImageResult<image::DynamicImage> {
    let mut reader = image::ImageReader::new(std::io::Cursor::new(data)).with_guessed_format()?;
    reader.limits(limits());
    reader.decode()
}

/// `photo` cut to a circle in the middle of a `width`×`height` canvas, which
/// is transparent around it, so the row's background shows there. The edge
/// fades out over a pixel to look round rather than jagged.
pub fn circle(photo: &DynamicImage, width: u32, height: u32) -> RgbaImage {
    let mut canvas = RgbaImage::new(width, height);
    let size = width.min(height);
    let side = photo.width().min(photo.height());
    if size == 0 || side == 0 {
        return canvas;
    }
    let square = photo
        .crop_imm(
            (photo.width() - side) / 2,
            (photo.height() - side) / 2,
            side,
            side,
        )
        .resize_exact(size, size, FilterType::Triangle)
        .to_rgba8();
    let (left, top) = ((width - size) / 2, (height - size) / 2);
    let radius = size as f32 / 2.0;
    for (x, y, pixel) in square.enumerate_pixels() {
        let dx = x as f32 + 0.5 - radius;
        let dy = y as f32 + 0.5 - radius;
        let inside = (radius - (dx * dx + dy * dy).sqrt() + 0.5).clamp(0.0, 1.0);
        let mut pixel = *pixel;
        pixel[3] = (f32::from(pixel[3]) * inside).round() as u8;
        canvas.put_pixel(left + x, top + y, pixel);
    }
    canvas
}

/// An image finished encoding (or failed) on a background thread.
pub struct ImageEvent {
    key: Key,
    /// [`Images::generation`] when it was started.
    generation: u64,
    result: Result<SlicedProtocol>,
}

enum FileState {
    Downloading,
    Ready(String),
    Failed,
}

/// An encoded chat photo.
struct Avatar {
    image: SlicedProtocol,
    /// When it was last drawn, in calls to [`Images::avatar`].
    used: u64,
    /// It has been drawn. Kitty's protocol sends an image to the terminal
    /// the first time it's drawn, and only then, so that time must not be
    /// under a popup that hides it.
    shown: bool,
}

pub struct Images {
    picker: Picker,
    tx: UnboundedSender<ImageEvent>,
    /// Encoded photos, with the frame each was last drawn in.
    ready: HashMap<Key, (SlicedProtocol, u64)>,
    /// Counts frames, for `ready`.
    frame: u64,
    /// Goes up when the files are forgotten, for a new TDLib client: an
    /// image still being built for the old one numbers its file the old way,
    /// and is dropped when it's done.
    generation: u64,
    /// Kept apart from `ready`: they stay when another chat is opened.
    avatars: HashMap<Key, Avatar>,
    /// Counts calls to [`Images::avatar`], to find the least recently drawn.
    avatar_clock: u64,
    /// Images being built, and when each started.
    building: HashMap<Key, std::time::Instant>,
    /// Keys that failed to decode; never retried.
    failed: HashSet<Key>,
    files: HashMap<i32, FileState>,
    /// Photos the last frame showed, with their size.
    wanted: Vec<(Preview, u16, u16)>,
    /// Chat photos the last frame showed, with their size.
    wanted_avatars: Vec<(ChatPhoto, u16, u16)>,
}

impl Images {
    pub fn new(picker: Picker, tx: UnboundedSender<ImageEvent>) -> Self {
        Self {
            picker,
            tx,
            ready: HashMap::new(),
            frame: 0,
            generation: 0,
            avatars: HashMap::new(),
            avatar_clock: 0,
            building: HashMap::new(),
            failed: HashSet::new(),
            files: HashMap::new(),
            wanted: Vec::new(),
            wanted_avatars: Vec::new(),
        }
    }

    /// False when the terminal shows no real images, only colored blocks
    /// (halfblocks), which are too coarse for a chat photo in a few cells.
    pub fn draws_photos(&self) -> bool {
        self.picker.protocol_type() != ProtocolType::Halfblocks
    }

    /// The terminal paints each image (or each of its rows) from one cell,
    /// and repaints the whole of it whenever that cell is sent again, over
    /// anything drawn on top since: sixel and iTerm2's protocol.
    pub fn paints_over(&self) -> bool {
        matches!(
            self.picker.protocol_type(),
            ProtocolType::Sixel | ProtocolType::Iterm2
        )
    }

    /// Pixel size of one terminal cell, for sizing photos by aspect ratio.
    pub fn font_size(&self) -> FontSize {
        self.picker.font_size()
    }

    /// The best image ready to draw: the photo, else its blurry thumbnail.
    pub fn get(&mut self, photo: &Preview, cols: u16, rows: u16) -> Option<&SlicedProtocol> {
        let key = |thumbnail| Key {
            file_id: photo.file_id,
            cols,
            rows,
            thumbnail,
            avatar: false,
        };
        let key = [key(false), key(true)]
            .into_iter()
            .find(|k| self.ready.contains_key(k))?;
        let frame = self.frame;
        let (image, used) = self.ready.get_mut(&key)?;
        *used = frame;
        Some(image)
    }

    /// True when the photo can't be shown: its download or decode failed.
    pub fn is_broken(&self, photo: &Preview) -> bool {
        matches!(self.files.get(&photo.file_id), Some(FileState::Failed))
    }

    /// Called while drawing, for each photo on screen.
    pub fn want(&mut self, photo: &Preview, cols: u16, rows: u16) {
        self.wanted.push((photo.clone(), cols, rows));
    }

    /// The best chat photo ready to draw, as for [`get`](Self::get). While
    /// `covered` (a popup may be over it), only one that has been drawn before
    /// is returned; see [`Avatar::shown`].
    pub fn avatar(
        &mut self,
        photo: &ChatPhoto,
        cols: u16,
        rows: u16,
        covered: bool,
    ) -> Option<&SlicedProtocol> {
        let key = |thumbnail| Key {
            file_id: photo.file_id,
            cols,
            rows,
            thumbnail,
            avatar: true,
        };
        // See paints_over: a popup may be over the list.
        if covered && self.paints_over() {
            return None;
        }
        let key = [key(false), key(true)]
            .into_iter()
            .find(|k| self.avatars.get(k).is_some_and(|a| a.shown || !covered))?;
        self.avatar_clock += 1;
        let avatar = self.avatars.get_mut(&key)?;
        avatar.used = self.avatar_clock;
        avatar.shown = true;
        Some(&avatar.image)
    }

    /// Called while drawing, for each chat photo on screen.
    pub fn want_avatar(&mut self, photo: &ChatPhoto, cols: u16, rows: u16) {
        self.wanted_avatars.push((photo.clone(), cols, rows));
    }

    /// Starts downloads and encodes for what the last frame wanted.
    pub fn fetch(&mut self, tg: &Tg) {
        self.frame += 1;
        self.fetch_avatars(tg);
        for (photo, cols, rows) in std::mem::take(&mut self.wanted) {
            let full = Key {
                file_id: photo.file_id,
                cols,
                rows,
                thumbnail: false,
                avatar: false,
            };
            if self.ready.contains_key(&full) {
                continue;
            }
            match self.files.get(&photo.file_id) {
                Some(FileState::Ready(path)) => {
                    let path = path.clone();
                    let limits = if photo.sticker {
                        sized(STICKER_SIDE)
                    } else {
                        limits()
                    };
                    self.build(full, move || Ok(open_within(&path, limits)?));
                }
                Some(FileState::Downloading | FileState::Failed) => {}
                None => {
                    self.files.insert(photo.file_id, FileState::Downloading);
                    tg.download(photo.file_id);
                }
            }
            if let Some(data) = photo.thumbnail {
                let key = Key {
                    thumbnail: true,
                    ..full
                };
                self.build(key, move || Ok(decode_bytes(&data)?));
            }
        }
    }

    fn fetch_avatars(&mut self, tg: &Tg) {
        let font = self.picker.font_size();
        for (photo, cols, rows) in std::mem::take(&mut self.wanted_avatars) {
            let full = Key {
                file_id: photo.file_id,
                cols,
                rows,
                thumbnail: false,
                avatar: true,
            };
            if self.avatars.contains_key(&full) {
                continue;
            }
            let (width, height) = (
                u32::from(cols) * u32::from(font.width),
                u32::from(rows) * u32::from(font.height),
            );
            let file = self
                .files
                .entry(photo.file_id)
                .or_insert_with(|| match &photo.path {
                    Some(path) => FileState::Ready(path.clone()),
                    None => {
                        tg.download_quiet(photo.file_id);
                        FileState::Downloading
                    }
                });
            if let FileState::Ready(path) = file {
                let path = path.clone();
                self.build(full, move || {
                    Ok(circle(&open_image(&path)?, width, height).into())
                });
            }
            if let Some(data) = photo.thumbnail {
                let key = Key {
                    thumbnail: true,
                    ..full
                };
                self.build(key, move || {
                    Ok(circle(&decode_bytes(&data)?, width, height).into())
                });
            }
        }
    }

    pub fn on_downloaded(&mut self, file_id: i32, path: Option<String>) {
        let state = path.map_or(FileState::Failed, FileState::Ready);
        self.files.insert(file_id, state);
    }

    pub fn on_built(&mut self, event: ImageEvent) {
        if event.generation != self.generation {
            return;
        }
        self.building.remove(&event.key);
        match event.result {
            Ok(image) if event.key.avatar => self.add_avatar(event.key, image),
            Ok(image) => self.add_ready(event.key, image),
            Err(_) => {
                self.failed.insert(event.key);
                if !event.key.thumbnail {
                    self.files.insert(event.key.file_id, FileState::Failed);
                }
            }
        }
    }

    fn add_ready(&mut self, key: Key, image: SlicedProtocol) {
        if self.ready.len() >= MAX_READY
            && let Some(oldest) = self
                .ready
                .iter()
                .min_by_key(|(_, (_, used))| *used)
                .map(|(&k, _)| k)
        {
            self.ready.remove(&oldest);
        }
        self.ready.insert(key, (image, self.frame));
    }

    fn add_avatar(&mut self, key: Key, image: SlicedProtocol) {
        if self.avatars.len() >= MAX_AVATARS
            && let Some(oldest) = self
                .avatars
                .iter()
                .min_by_key(|(_, a)| a.used)
                .map(|(&k, _)| k)
        {
            self.avatars.remove(&oldest);
        }
        let avatar = Avatar {
            image,
            used: self.avatar_clock,
            shown: false,
        };
        self.avatars.insert(key, avatar);
    }

    /// Drops encoded photos, e.g. when switching chats. Chat photos and
    /// downloads stay.
    pub fn clear(&mut self) {
        self.ready.clear();
    }

    /// Forgets every file, for a new TDLib client: it numbers files afresh.
    pub fn forget_files(&mut self) {
        self.generation += 1;
        self.ready.clear();
        self.avatars.clear();
        self.building.clear();
        self.failed.clear();
        self.files.clear();
        self.wanted.clear();
        self.wanted_avatars.clear();
    }

    /// Decodes and encodes on a blocking thread, once per key.
    fn build(
        &mut self,
        key: Key,
        decode: impl FnOnce() -> Result<image::DynamicImage> + Send + 'static,
    ) {
        let built = if key.avatar {
            self.avatars.contains_key(&key)
        } else {
            self.ready.contains_key(&key)
        };
        // One that never finishes mustn't hold its place for good: it counts
        // as failed, and the others get their turn.
        let failed = &mut self.failed;
        self.building.retain(|&key, started| {
            let alive = started.elapsed() < BUILD_TIMEOUT;
            if !alive {
                failed.insert(key);
            }
            alive
        });
        if built
            || self.building.contains_key(&key)
            || self.failed.contains(&key)
            || self.building.len() >= MAX_BUILDING
        {
            return;
        }
        self.building.insert(key, std::time::Instant::now());
        let picker = self.picker.clone();
        let tx = self.tx.clone();
        let generation = self.generation;
        tokio::task::spawn_blocking(move || {
            let result = contained(|| {
                decode().and_then(|image| {
                    let size = Size::new(key.cols, key.rows);
                    Ok(SlicedProtocol::new_with_resize(
                        &picker, image, size, RESIZE,
                    )?)
                })
            });
            let _ = tx.send(ImageEvent {
                key,
                generation,
                result,
            });
        });
    }

    #[cfg(test)]
    pub fn insert_ready(&mut self, key: Key, image: image::DynamicImage) {
        let size = Size::new(key.cols, key.rows);
        let image = SlicedProtocol::new_with_resize(&self.picker, image, size, RESIZE).unwrap();
        if key.avatar {
            self.add_avatar(key, image);
        } else {
            self.add_ready(key, image);
        }
    }
}

/// Shared thumbnail bytes, so cloning a [`Preview`] each frame is cheap.
pub type Thumbnail = Arc<[u8]>;

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

    fn png(width: u32, height: u32) -> Vec<u8> {
        let mut data = std::io::Cursor::new(Vec::new());
        image::RgbImage::new(width, height)
            .write_to(&mut data, image::ImageFormat::Png)
            .unwrap();
        data.into_inner()
    }

    #[test]
    fn chat_photos_are_cut_to_a_circle_with_see_through_corners() {
        let photo = image::RgbaImage::from_pixel(160, 100, image::Rgba([255, 0, 0, 255])).into();
        let round = circle(&photo, 40, 44);
        let alpha = |x, y| round.get_pixel(x, y)[3];
        assert_eq!(alpha(0, 2), 0, "corner");
        assert_eq!(alpha(39, 41), 0, "corner");
        assert_eq!(alpha(20, 22), 255, "middle");
        assert_eq!(alpha(20, 0), 0, "above the circle: the canvas is taller");
        assert_eq!(alpha(1, 22), 255, "the circle reaches the sides");
        let edge = alpha(6, 7);
        assert!(edge > 0 && edge < 255, "a soft edge: {edge}");
        assert_eq!(circle(&photo, 0, 44).dimensions(), (0, 44), "no panic");
    }

    fn key(file_id: i32) -> Key {
        Key {
            file_id,
            cols: 4,
            rows: 2,
            thumbnail: false,
            avatar: false,
        }
    }

    fn preview(file_id: i32) -> Preview {
        Preview {
            file_id,
            width: 8,
            height: 8,
            thumbnail: None,
            sticker: false,
        }
    }

    #[tokio::test]
    async fn a_decoder_panic_marks_the_photo_broken_instead_of_ending_the_app() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        let mut images = Images::new(Picker::halfblocks(), tx);
        images.build(key(7), || panic!("a decoder bug"));
        images.on_built(rx.recv().await.unwrap());
        assert!(images.is_broken(&preview(7)));
        assert!(!panic_is_contained(), "only while decoding");
    }

    #[tokio::test]
    async fn an_image_built_for_a_logged_out_session_is_dropped() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        let mut images = Images::new(Picker::halfblocks(), tx);
        let (go, wait) = std::sync::mpsc::channel::<()>();
        images.build(key(7), move || {
            wait.recv().ok();
            Ok(DynamicImage::new_rgba8(8, 8))
        });
        images.forget_files();
        go.send(()).unwrap();
        images.on_built(rx.recv().await.unwrap());
        assert!(
            images.get(&preview(7), 4, 2).is_none(),
            "not the new session's"
        );
        assert!(!images.is_broken(&preview(7)));
    }

    #[tokio::test]
    async fn only_a_few_images_decode_at_once() {
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let mut images = Images::new(Picker::halfblocks(), tx);
        let (go, wait) = std::sync::mpsc::channel::<()>();
        let wait = std::sync::Arc::new(std::sync::Mutex::new(wait));
        for id in 0..10 {
            let wait = wait.clone();
            images.build(key(id), move || {
                let _ = wait.lock().unwrap().recv();
                Ok(DynamicImage::new_rgba8(8, 8))
            });
        }
        assert_eq!(images.building.len(), MAX_BUILDING);
        drop(go);
    }

    #[tokio::test]
    async fn a_decode_that_never_finishes_gives_up_its_place() {
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let mut images = Images::new(Picker::halfblocks(), tx);
        for id in 0..MAX_BUILDING as i32 {
            let long_ago = std::time::Instant::now() - BUILD_TIMEOUT;
            images.building.insert(key(id), long_ago);
        }
        images.build(key(99), || Ok(DynamicImage::new_rgba8(8, 8)));
        assert!(images.building.contains_key(&key(99)), "its turn came");
        assert!(
            images.failed.contains(&key(0)),
            "the stuck one counts as failed"
        );
    }

    #[test]
    fn the_photo_cache_keeps_the_ones_drawn_lately() {
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let mut images = Images::new(Picker::halfblocks(), tx);
        images.insert_ready(key(0), DynamicImage::new_rgba8(8, 8));
        for id in 1..MAX_READY as i32 + 10 {
            images.frame += 1;
            assert!(images.get(&preview(0), 4, 2).is_some(), "still drawn");
            images.insert_ready(key(id), DynamicImage::new_rgba8(8, 8));
        }
        assert_eq!(images.ready.len(), MAX_READY);
        assert!(images.get(&preview(0), 4, 2).is_some());
        assert!(images.get(&preview(1), 4, 2).is_none(), "the oldest went");
    }

    #[test]
    fn images_bigger_than_telegram_sends_are_refused() {
        assert!(decode_bytes(&png(512, 512)).is_ok());
        let error = decode_bytes(&png(5000, 1)).unwrap_err();
        assert!(matches!(error, image::ImageError::Limits(_)), "{error}");
    }
}