//! 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));
    }
}