bao-browser 0.4.0

// @trace REQ-BRW-002  REQ-CDP-007: Screenshot capture via servo rendering pipeline
use std::io::Cursor;

use image::{ImageFormat, RgbaImage};

use crate::error::BrowserError;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScreenshotFormat {
    Png,
    Jpeg,
    WebP,
}

pub fn encode_image(image: &RgbaImage, format: ScreenshotFormat) -> Result<Vec<u8>, BrowserError> {
    encode_image_with_quality(image, format, DEFAULT_JPEG_QUALITY)
}

/// Default JPEG quality for the plain `encode_image` face (screencast
/// carries its own `quality` param per the CDP spec, Chrome default 80).
const DEFAULT_JPEG_QUALITY: u8 = 80;

/// Encode with an explicit JPEG quality (1..=100, CDP `Page.captureScreenshot`
/// / `Page.startScreencast` `quality` param — REQ-CDP-009). PNG and WebP
/// ignore the quality knob (lossless / fixed-quality codecs).
// @trace REQ-CDP-009 [criterion:format 参数支持 jpeg 与 png]
pub fn encode_image_with_quality(
    image: &RgbaImage,
    format: ScreenshotFormat,
    jpeg_quality: u8,
) -> Result<Vec<u8>, BrowserError> {
    let mut buf = Cursor::new(Vec::with_capacity(
        image.width() as usize * image.height() as usize,
    ));
    match format {
        ScreenshotFormat::Png => image
            .write_to(&mut buf, ImageFormat::Png)
            .map_err(|e| BrowserError::Rendering(format!("PNG encode failed: {e}")))?,
        ScreenshotFormat::Jpeg => {
            let rgb = image::DynamicImage::ImageRgba8(image.clone()).to_rgb8();
            let encoder =
                image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, jpeg_quality);
            rgb.write_with_encoder(encoder)
                .map_err(|e| BrowserError::Rendering(format!("JPEG encode failed: {e}")))?;
        }
        ScreenshotFormat::WebP => image
            .write_to(&mut buf, ImageFormat::WebP)
            .map_err(|e| BrowserError::Rendering(format!("WebP encode failed: {e}")))?,
    }
    Ok(buf.into_inner())
}

#[cfg(test)]
mod tests {
    // @trace REQ-BRW-002 [req:REQ-BRW-002,REQ-CDP-007] [level:unit]
    use super::*;
    use image::Rgba;

    fn red_image(w: u32, h: u32) -> RgbaImage {
        RgbaImage::from_pixel(w, h, Rgba([255, 0, 0, 255]))
    }

    /// Encode a solid-red WxH image in `fmt` and assert the output is
    /// non-empty and begins with `magic` (the container's magic bytes) —
    /// the shared shape of every encode smoke test below.
    fn assert_encode_magic(w: u32, h: u32, fmt: ScreenshotFormat, magic: &[u8]) {
        let img = red_image(w, h);
        let out = encode_image(&img, fmt).unwrap();
        assert!(!out.is_empty());
        assert_eq!(&out[0..magic.len()], magic);
    }

    #[test]
    fn encode_1x1_png_returns_nonempty_with_magic() {
        assert_encode_magic(1, 1, ScreenshotFormat::Png, &[0x89, 0x50, 0x4E, 0x47]);
    }

    #[test]
    fn encode_1x1_jpeg_returns_nonempty_with_magic() {
        assert_encode_magic(1, 1, ScreenshotFormat::Jpeg, &[0xFF, 0xD8]);
    }

    #[test]
    fn encode_100x100_png_produces_valid_data() {
        assert_encode_magic(100, 100, ScreenshotFormat::Png, &[0x89, 0x50, 0x4E, 0x47]);
    }

    #[test]
    fn encode_100x100_jpeg_produces_valid_data() {
        assert_encode_magic(100, 100, ScreenshotFormat::Jpeg, &[0xFF, 0xD8]);
    }

    #[test]
    fn encode_1x1_webp_returns_nonempty_with_riff_magic() {
        // WebP files start with "RIFF" header
        assert_encode_magic(1, 1, ScreenshotFormat::WebP, b"RIFF");
    }

    #[test]
    fn encode_100x100_webp_produces_valid_data() {
        assert_encode_magic(100, 100, ScreenshotFormat::WebP, b"RIFF");
    }

    #[test]
    fn encode_empty_image_no_panic() {
        let img = red_image(0, 0);
        let _ = encode_image(&img, ScreenshotFormat::Png);
        let _ = encode_image(&img, ScreenshotFormat::Jpeg);
        let _ = encode_image(&img, ScreenshotFormat::WebP);
    }

    // @trace REQ-CDP-009 [criterion:format 参数支持 jpeg 与 png]
    #[test]
    fn encode_with_quality_both_formats_valid_magic() {
        let img = red_image(64, 64);
        let jpeg = encode_image_with_quality(&img, ScreenshotFormat::Jpeg, 60).unwrap();
        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
        let png = encode_image_with_quality(&img, ScreenshotFormat::Png, 60).unwrap();
        assert_eq!(&png[0..4], &[0x89, 0x50, 0x4E, 0x47]);
    }

    // @trace REQ-CDP-009 [criterion:format 参数支持 jpeg 与 png]
    #[test]
    fn encode_jpeg_lower_quality_is_not_larger() {
        let img = red_image(200, 200);
        // A noise-free solid frame can compress to identical size at both
        // qualities, but a LOWER quality must never inflate the output.
        let hi = encode_image_with_quality(&img, ScreenshotFormat::Jpeg, 95)
            .unwrap()
            .len();
        let lo = encode_image_with_quality(&img, ScreenshotFormat::Jpeg, 20)
            .unwrap()
            .len();
        assert!(
            lo <= hi,
            "quality 20 ({lo}B) must not exceed quality 95 ({hi}B)"
        );
    }

    #[test]
    fn all_formats_produce_nontrivial_output() {
        let img = red_image(200, 200);
        let png = encode_image(&img, ScreenshotFormat::Png).unwrap();
        let jpeg = encode_image(&img, ScreenshotFormat::Jpeg).unwrap();
        let webp = encode_image(&img, ScreenshotFormat::WebP).unwrap();
        assert!(
            png.len() > 100,
            "PNG should be nontrivial: {} bytes",
            png.len()
        );
        assert!(
            jpeg.len() > 100,
            "JPEG should be nontrivial: {} bytes",
            jpeg.len()
        );
        assert!(
            webp.len() > 100,
            "WebP should be nontrivial: {} bytes",
            webp.len()
        );
    }
}