//! Provides basic support for rendering dumps of X11 windows made with [the `xwud` utility](https://gitlab.freedesktop.org/xorg/app/xwud/-/blob/master/xwud.c?ref_type=heads).

use std::{
    ffi::CString,
    io::Cursor,
    mem,
};

const SUPPORTED_XWD_VERSION: u32 = 7;

/// Header for `xwd`-style dumps of X11 windows.
///
/// This is taken from `/usr/include/X11/XWDFile.h` on a Debian system.
#[derive(Clone)]
#[repr(C)]
pub struct XwdHeader {
    /// Total header size (including null-terminated window name).
    pub size: u32,
    /// Version of `xwd` that this dump came from.
    pub version: u32,
    pub _format: u32,
    /// Bit depth of the X11 server.
    pub depth: u32,
    /// Image width.
    pub width: u32,
    /// Image height.
    pub height: u32,
    pub _x_offset: u32,
    pub _byte_order: u32,
    pub _bitmap_unit: u32,
    pub _bitmap_bit_order: u32,
    pub _bitmap_pad: u32,
    pub bits_per_pixel: u32,
    pub _bytes_per_line: u32,
    /// X11 visual class for this window. See [`XwdImage::visual_class`] for a
    /// more human-interpretable value.
    pub visual_class: u32,
    /// Bitmask for the red channel in pixel data.
    pub red_mask: u32,
    /// Bitmask for the green channel in pixel data.
    pub green_mask: u32,
    /// Bitmask for the blue channel in pixel data.
    pub blue_mask: u32,
    /// Number of bits per RGB tuple in color data. (E.g., 24-bit color may
    /// actually be padded to 4 bytes.)
    pub bits_per_rgb: u32,
    pub _colormap_entries: u32,
    /// Number of colors in the colormap for the window that this dump came from.
    ///
    /// Note that a colormap may exist but be unused (if the visual class is not
    /// one that uses a colormap).
    pub ncolors: u32,
    pub _window_width: u32,
    pub _window_height: u32,
    pub _window_x: u32,
    pub _window_y: u32,
    pub _window_border_width: u32,
}

const HEADER_SIZE: u32 = mem::size_of::<XwdHeader>() as u32;

/// Color data for an entry in a colormap.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(C)]
pub struct XwdColor {
    pixel: u32,
    red: u16,
    green: u16,
    blue: u16,
    flags: u8,
    pad: u8,
}

fn next_u32(bytes: &mut impl Iterator<Item = u8>) -> u32 {
    let bytes = [
        bytes.next().unwrap(),
        bytes.next().unwrap(),
        bytes.next().unwrap(),
        bytes.next().unwrap(),
    ];
    u32::from_be_bytes(bytes)
}

fn next_u16(bytes: &mut impl Iterator<Item = u8>) -> u16 {
    let bytes = [bytes.next().unwrap(), bytes.next().unwrap()];
    u16::from_be_bytes(bytes)
}

/// Describes how pixel data is represented in an `xwd` dump.
///
/// See descriptions of visual classes at
/// <https://tronche.com/gui/x/xlib/window/visual-types.html> and [the official
/// documentation](https://www.x.org/releases/X11R7.7/doc/xproto/x11protocol.html#visual_information)
/// for more information.
#[derive(Clone, Copy, Eq, Debug, PartialEq)]
pub enum VisualClass {
    StaticGray,
    GrayScale,
    StaticColor,
    PseudoColor,
    TrueColor,
    DirectColor,
}

impl XwdHeader {
    /// Decodes an `XwdHeader` from `bytes`, panicking on failure.
    ///
    /// Leaves `bytes` pointing to the start of the window name.
    pub fn read(bytes: &mut impl Iterator<Item = u8>) -> XwdHeader {
        // This layout might change in the future, but let's just go with it for now.
        XwdHeader {
            size: next_u32(bytes),
            version: next_u32(bytes),
            _format: next_u32(bytes),
            depth: next_u32(bytes),
            width: next_u32(bytes),
            height: next_u32(bytes),
            _x_offset: next_u32(bytes),
            _byte_order: next_u32(bytes),
            _bitmap_unit: next_u32(bytes),
            _bitmap_bit_order: next_u32(bytes),
            _bitmap_pad: next_u32(bytes),
            bits_per_pixel: next_u32(bytes),
            _bytes_per_line: next_u32(bytes),
            visual_class: next_u32(bytes),
            red_mask: next_u32(bytes),
            green_mask: next_u32(bytes),
            blue_mask: next_u32(bytes),
            bits_per_rgb: next_u32(bytes),
            _colormap_entries: next_u32(bytes),
            ncolors: next_u32(bytes),
            _window_width: next_u32(bytes),
            _window_height: next_u32(bytes),
            _window_x: next_u32(bytes),
            _window_y: next_u32(bytes),
            _window_border_width: next_u32(bytes),
        }
    }
}

impl XwdColor {
    /// Reads an `XwdColor` from `bytes`, panicking on failure.
    ///
    /// Leaves `bytes` pointing to the next byte after the end of the color
    /// definition.
    pub fn read(bytes: &mut impl Iterator<Item = u8>) -> XwdColor {
        XwdColor {
            pixel: next_u32(bytes),
            red: next_u16(bytes),
            green: next_u16(bytes),
            blue: next_u16(bytes),
            flags: bytes.next().unwrap(),
            pad: bytes.next().unwrap(),
        }
    }
}

/// A decoded `xwd` dump.
pub struct XwdImage {
    /// Header with image metadata.
    pub header: XwdHeader,
    /// Name of the X11 window that the image came from.
    pub name: CString,
    /// Colormap data for the X11 window that the image came from.
    ///
    /// Note that this may be populated even if it is unused (because the
    /// window's visual class is one that does not use a colormap).
    pub colors: Vec<XwdColor>,
    /// Pixel data for the image.
    ///
    /// This may be indices pointing into `colors`, or it may be actual pixel
    /// values.
    pub pixels: Vec<u8>,
}

impl XwdImage {
    /// Reads an `XwdImage` from `bytes`, panicking on failure.
    pub fn read(bytes: &mut impl Iterator<Item = u8>) -> XwdImage {
        let header = XwdHeader::read(bytes);
        if header.version != SUPPORTED_XWD_VERSION {
            panic!(
                "header version {} is not supported. only version {} files are supported.",
                header.version, SUPPORTED_XWD_VERSION
            );
        }
        let name = Self::read_name(&header, bytes);
        let colors = Self::read_colors(&header, bytes);
        let pixels = bytes.collect();

        XwdImage {
            header,
            name,
            colors,
            pixels,
        }
    }

    /// Reads a colormap for the image described by `header` from `bytes`.
    ///
    /// `bytes` must point at the start of the colormap data. Leaves `bytes`
    /// pointing at the first byte of the next item in the data stream.
    pub fn read_colors(header: &XwdHeader, bytes: &mut impl Iterator<Item = u8>) -> Vec<XwdColor> {
        let mut colors = Vec::with_capacity(header.ncolors as usize);
        for _ in 0..header.ncolors {
            colors.push(XwdColor::read(bytes));
        }
        colors
    }

    /// Reads the X11 window name from `bytes`.
    ///
    /// `bytes` must point at the start of the window name. Leaves `bytes`
    /// pointing at the first byte of the next item in the data stream.
    ///
    /// Note that the window name is a null-terminated string, but `header.size`
    /// (which is supposed to be equal to the size of the header structure plus
    /// the window name) may actually be longer than `strlen(name) +
    /// sizeof(HeaderType)`. When that is the case, the extra padding after the
    /// first null byte in the window name will be passed over silently.
    pub fn read_name(header: &XwdHeader, bytes: &mut impl Iterator<Item = u8>) -> CString {
        if header.size <= HEADER_SIZE {
            panic!(
                "Invalid header size {} (smaller than header struct size {})",
                header.size, HEADER_SIZE
            );
        }
        let mut buf = Vec::new();
        let mut found_nul = false;
        for _ in 0..header.size - HEADER_SIZE {
            let b = bytes.next().unwrap();
            if b == 0 {
                if found_nul {
                    continue;
                } else {
                    buf.push(b);
                    found_nul = true;
                }
            } else {
                buf.push(b);
            }
        }
        CString::from_vec_with_nul(buf).unwrap()
    }

    /// Returns the visual class for this window dump.
    pub fn visual_class(&self) -> VisualClass {
        match self.header.visual_class {
            0 => VisualClass::StaticGray,
            1 => VisualClass::GrayScale,
            2 => VisualClass::StaticColor,
            3 => VisualClass::PseudoColor,
            4 => VisualClass::TrueColor,
            5 => VisualClass::DirectColor,
            _ => panic!(
                "unrecognized X11 visual class: {}",
                self.header.visual_class
            ),
        }
    }

    /// Encodes this image as a PNG, consumes `self` in the process (because the
    /// image data is reused during the encoding process).
    ///
    /// Note that only [`VisualClass::TrueColor`] images with a specific pixel
    /// format (24-bit color packed into 32-bit segments with BGR ordering) are
    /// supported.
    pub fn into_png(self) -> Result<Vec<u8>, png::EncodingError> {
        match self.visual_class() {
            VisualClass::TrueColor => into_png_true_color(self.header, self.pixels),
            _ => todo!("X11 visual classes other than TrueColor are not yet supported"),
        }
    }
}

fn into_png_true_color(
    header: XwdHeader,
    mut pixels: Vec<u8>,
) -> Result<Vec<u8>, png::EncodingError> {
    if header.depth != 24 {
        todo!("pixmap color depths other than 24-bit are not yet supported");
    }
    if !(header.red_mask == 0xff0000
        && header.green_mask == 0x00ff00
        && header.blue_mask == 0x0000ff)
    {
        todo!("color orderings other than RGB are not yet supported");
    }
    if header.bits_per_rgb == 32 {
        todo!("bits per rgb triplet other than 32 are not yet supported");
    }

    let mut out = Cursor::new(Vec::<u8>::new());
    let mut encoder = png::Encoder::new(&mut out, header.width, header.height);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    encoder.set_compression(png::Compression::High);

    for pixel in pixels.chunks_mut(4) {
        let blue = pixel[0];
        let green = pixel[1];
        let red = pixel[2];
        pixel[0] = red;
        pixel[1] = green;
        pixel[2] = blue;
        pixel[3] = 0xFF;
    }

    let mut writer = encoder.write_header()?;
    writer.write_image_data(&pixels)?;
    writer.finish()?;

    Ok(out.into_inner())
}

#[cfg(test)]
mod tests {
    use super::XwdImage;
    use insta_image::assert_png_snapshot;

    #[test]
    fn png_encode() {
        let snowlamp_xwd = XwdImage::read(&mut include_bytes!("snowlamp.xwd").into_iter().copied());
        assert_png_snapshot!("snowlamp_encoded", snowlamp_xwd.into_png().unwrap());
    }
}
