Pup - the AI's own browser

//! PNG fit-inside resize without the `image` crate: decode with `png`, box-filter downscale on
//! RGBA, re-encode. Replaces the Node build's sharp dependency for screenshots (Node: processScreenshot).

use std::io::Cursor;

/// Returns (png bytes, width, height, origWidth, origHeight, resized). Images already within
/// `max_dim` on both edges are returned re-encoded only when they needed expansion, else as-is.
pub fn fit_png(bytes: &[u8], max_dim: u32) -> Result<(Vec<u8>, u32, u32, u32, u32, bool), String> {
    let mut dec = png::Decoder::new(Cursor::new(bytes));
    dec.set_transformations(png::Transformations::EXPAND | png::Transformations::STRIP_16 | png::Transformations::ALPHA);
    let mut reader = dec.read_info().map_err(|e| format!("png decode: {e}"))?;
    let info = reader.info().clone();
    let (ow, oh) = (info.width, info.height);
    if ow <= max_dim && oh <= max_dim { return Ok((bytes.to_vec(), ow, oh, ow, oh, false)); }
    let mut buf = vec![0u8; reader.output_buffer_size()];
    let out = reader.next_frame(&mut buf).map_err(|e| format!("png frame: {e}"))?;
    let bpp = match out.color_type { png::ColorType::Rgba => 4, png::ColorType::GrayscaleAlpha => 2, png::ColorType::Rgb => 3, png::ColorType::Grayscale => 1, _ => return Err("unsupported png colour type".into()) };
    let rgba: Vec<u8> = match bpp {
        4 => buf[..out.buffer_size()].to_vec(),
        3 => buf[..out.buffer_size()].chunks(3).flat_map(|p| [p[0], p[1], p[2], 255]).collect(),
        2 => buf[..out.buffer_size()].chunks(2).flat_map(|p| [p[0], p[0], p[0], p[1]]).collect(),
        _ => buf[..out.buffer_size()].iter().flat_map(|&g| [g, g, g, 255]).collect(),
    };
    let scale = (max_dim as f64 / ow as f64).min(max_dim as f64 / oh as f64);
    let nw = ((ow as f64 * scale).round() as u32).max(1);
    let nh = ((oh as f64 * scale).round() as u32).max(1);
    let small = downscale_rgba(&rgba, ow, oh, nw, nh);
    let mut out_bytes = Vec::new();
    {
        let mut enc = png::Encoder::new(&mut out_bytes, nw, nh);
        enc.set_color(png::ColorType::Rgba); enc.set_depth(png::BitDepth::Eight);
        enc.set_compression(png::Compression::Best); enc.set_filter(png::FilterType::Paeth);
        let mut w = enc.write_header().map_err(|e| format!("png encode: {e}"))?;
        w.write_image_data(&small).map_err(|e| format!("png write: {e}"))?;
    }
    Ok((out_bytes, nw, nh, ow, oh, true))
}

/// Area-averaging downscale (each output pixel averages the source box it covers).
pub fn downscale_rgba(src: &[u8], sw: u32, sh: u32, dw: u32, dh: u32) -> Vec<u8> {
    let mut out = vec![0u8; (dw * dh * 4) as usize];
    let xs = sw as f64 / dw as f64; let ys = sh as f64 / dh as f64;
    for dy in 0..dh {
        let y0 = (dy as f64 * ys).floor() as u32; let y1 = (((dy + 1) as f64 * ys).ceil() as u32).min(sh).max(y0 + 1);
        for dx in 0..dw {
            let x0 = (dx as f64 * xs).floor() as u32; let x1 = (((dx + 1) as f64 * xs).ceil() as u32).min(sw).max(x0 + 1);
            let mut acc = [0u64; 4]; let mut n = 0u64;
            for y in y0..y1 { let row = (y * sw) as usize * 4; for x in x0..x1 { let i = row + (x as usize) * 4; for c in 0..4 { acc[c] += src[i + c] as u64; } n += 1; } }
            let o = ((dy * dw + dx) * 4) as usize;
            for c in 0..4 { out[o + c] = (acc[c] / n.max(1)) as u8; }
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    fn solid_png(w: u32, h: u32) -> Vec<u8> {
        let mut v = Vec::new();
        { let mut e = png::Encoder::new(&mut v, w, h); e.set_color(png::ColorType::Rgb); e.set_depth(png::BitDepth::Eight); let mut wr = e.write_header().unwrap(); wr.write_image_data(&vec![120u8; (w * h * 3) as usize]).unwrap(); }
        v
    }
    #[test] fn fits_inside() {
        let p = solid_png(3000, 1000);
        let (out, w, h, ow, oh, resized) = fit_png(&p, 1568).unwrap();
        assert!(resized); assert_eq!((ow, oh), (3000, 1000)); assert_eq!(w, 1568); assert_eq!(h, 523); assert!(!out.is_empty());
        let (_, w2, h2, _, _, r2) = fit_png(&solid_png(800, 600), 1568).unwrap(); assert!(!r2); assert_eq!((w2, h2), (800, 600));
        let (o2, ..) = fit_png(&out, 1568).unwrap(); assert_eq!(o2.len(), out.len());
    }
}