use anyhow::{Result, anyhow};
use image::ImageEncoder;
use std::path::Path;
use x11rb::connection::Connection;
use x11rb::protocol::xfixes::ConnectionExt as _;
use x11rb::protocol::xproto::{ConnectionExt as _, ImageFormat, ImageOrder};
use x11rb::rust_connection::RustConnection;

pub struct Image {
    pub w: usize,
    pub h: usize,
    pub px: Vec<u8>,
}

impl Image {
    #[inline]
    pub fn at(&self, x: usize, y: usize) -> (u8, u8, u8) {
        let o = (y * self.w + x) * 3;
        (self.px[o], self.px[o + 1], self.px[o + 2])
    }

    #[cfg(test)]
    pub fn decode_png(bytes: &[u8]) -> Result<Image> {
        let img = image::load_from_memory(bytes)?.to_rgb8();
        let (w, h) = img.dimensions();
        Ok(Image {
            w: w as usize,
            h: h as usize,
            px: img.into_raw(),
        })
    }

    pub fn to_png(&self) -> Result<Vec<u8>> {
        encode_png(&self.px, self.w, self.h)
    }
}

pub fn grab(display: &str) -> Option<Vec<u8>> {
    (|| -> Result<Vec<u8>> {
        let (conn, mut img) = capture(display)?;
        composite_cursor(&conn, &mut img);
        img.to_png()
    })()
    .map_err(|e| tracing::debug!("frame grab failed: {e}"))
    .ok()
}

pub fn grab_image(display: &str) -> Result<Image> {
    Ok(capture(display)?.1)
}

pub fn screenshot(display: &str, path: &Path) -> Result<()> {
    std::fs::write(path, capture(display)?.1.to_png()?)?;
    Ok(())
}

fn capture(display: &str) -> Result<(RustConnection, Image)> {
    let (conn, screen_num) = x11rb::connect(Some(display))?;
    let screen = &conn.setup().roots[screen_num];
    let root = screen.root;
    let (w, h) = (
        screen.width_in_pixels as usize,
        screen.height_in_pixels as usize,
    );
    let img = conn
        .get_image(ImageFormat::Z_PIXMAP, root, 0, 0, w as u16, h as u16, !0u32)?
        .reply()?;
    let bpp = conn
        .setup()
        .pixmap_formats
        .iter()
        .find(|f| f.depth == img.depth)
        .map(|f| f.bits_per_pixel as usize)
        .ok_or_else(|| anyhow!("no pixmap format for depth {}", img.depth))?;
    if bpp != 32 {
        return Err(anyhow!("unexpected {bpp} bits/pixel (want 32)"));
    }
    let lsb = conn.setup().image_byte_order == ImageOrder::LSB_FIRST;
    let px = zpixmap_to_rgb(&img.data, w, h, lsb);
    Ok((conn, Image { w, h, px }))
}

fn composite_cursor(conn: &RustConnection, img: &mut Image) {
    let cursor = (|| {
        conn.xfixes_query_version(5, 0).ok()?.reply().ok()?;
        conn.xfixes_get_cursor_image().ok()?.reply().ok()
    })();
    let Some(cur) = cursor else { return };
    let (w, h) = (img.w, img.h);
    let (cw, ch) = (cur.width as i32, cur.height as i32);
    let (ox, oy) = (
        cur.x as i32 - cur.xhot as i32,
        cur.y as i32 - cur.yhot as i32,
    );
    for cy in 0..ch {
        for cx in 0..cw {
            let (px, py) = (ox + cx, oy + cy);
            if px < 0 || py < 0 || px >= w as i32 || py >= h as i32 {
                continue;
            }
            let argb = cur.cursor_image[(cy * cw + cx) as usize];
            let a = (argb >> 24) & 0xff;
            if a == 0 {
                continue;
            }
            let (cr, cg, cb) = ((argb >> 16) & 0xff, (argb >> 8) & 0xff, argb & 0xff);
            let inv = 255 - a;
            let p = (py as usize * w + px as usize) * 3;
            img.px[p] = (cr + img.px[p] as u32 * inv / 255) as u8;
            img.px[p + 1] = (cg + img.px[p + 1] as u32 * inv / 255) as u8;
            img.px[p + 2] = (cb + img.px[p + 2] as u32 * inv / 255) as u8;
        }
    }
}

fn encode_png(rgb: &[u8], w: usize, h: usize) -> Result<Vec<u8>> {
    let mut out = Vec::new();
    image::codecs::png::PngEncoder::new(&mut out).write_image(
        rgb,
        w as u32,
        h as u32,
        image::ExtendedColorType::Rgb8,
    )?;
    Ok(out)
}

fn zpixmap_to_rgb(data: &[u8], w: usize, h: usize, lsb: bool) -> Vec<u8> {
    let stride = w * 4;
    let mut rgb = vec![0u8; w * h * 3];
    for y in 0..h {
        for x in 0..w {
            let o = y * stride + x * 4;
            let (r, g, b) = if lsb {
                (data[o + 2], data[o + 1], data[o])
            } else {
                (data[o + 1], data[o + 2], data[o + 3])
            };
            let p = (y * w + x) * 3;
            rgb[p] = r;
            rgb[p + 1] = g;
            rgb[p + 2] = b;
        }
    }
    rgb
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::solver::display::VirtualDisplay;

    #[test]
    fn zpixmap_color_order() {
        let rgb = zpixmap_to_rgb(&[0, 0, 255, 0], 1, 1, true);
        assert_eq!(&rgb, &[255, 0, 0], "LSB red mis-ordered");
        let rgb = zpixmap_to_rgb(&[0, 255, 0, 0], 1, 1, false);
        assert_eq!(&rgb, &[255, 0, 0], "MSB red mis-ordered");
    }

    #[test]
    #[ignore = "needs Xvfb; run with --ignored"]
    fn composites_cursor() {
        use x11rb::protocol::xtest::ConnectionExt as _;
        let d = VirtualDisplay::start(800, 600).expect("start Xvfb");
        let name = d.name();
        std::thread::sleep(std::time::Duration::from_millis(300));
        let (conn, sn) = x11rb::connect(Some(&name)).expect("connect");
        let root = conn.setup().roots[sn].root;
        conn.xtest_fake_input(
            x11rb::protocol::xproto::MOTION_NOTIFY_EVENT,
            0,
            0,
            root,
            400,
            300,
            0,
        )
        .and_then(|_| conn.flush())
        .expect("warp cursor");
        std::thread::sleep(std::time::Duration::from_millis(200));

        let png = grab(&name).expect("grab a frame");
        let img = image::load_from_memory(&png).expect("valid PNG").to_rgb8();
        assert_eq!(img.dimensions(), (800, 600));

        let cursor_pixels = (380..420)
            .flat_map(|x| (290..330).map(move |y| (x, y)))
            .filter(|&(x, y)| img.get_pixel(x, y).0 != [0, 0, 0])
            .count();
        assert!(
            cursor_pixels > 10,
            "cursor not composited near (400,300): {cursor_pixels} px"
        );
    }
}
