use anyhow::{anyhow, Result};
use std::fs;
#[derive(Debug, Clone)]
pub struct ImageInfo {
pub format: ImageFormat,
pub width: u32,
pub height: u32,
pub data: Vec<u8>,
pub bits_per_component: u8,
pub color_components: u8, pub alt_text: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ImageFormat {
Jpeg,
Png,
Bmp,
}
pub fn detect_image_format(data: &[u8]) -> Result<ImageFormat> {
if data.len() < 4 {
return Err(anyhow!("Image data too short"));
}
if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
Ok(ImageFormat::Jpeg)
} else if data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 {
Ok(ImageFormat::Png)
} else if data[0] == 0x42 && data[1] == 0x4D {
Ok(ImageFormat::Bmp)
} else {
Err(anyhow!("Unsupported image format"))
}
}
pub fn load_image(path: &str) -> Result<ImageInfo> {
load_image_with_alt_text(path, None)
}
pub fn load_image_with_alt_text(path: &str, alt_text: Option<String>) -> Result<ImageInfo> {
let data = fs::read(path)?;
let format = detect_image_format(&data)?;
let (width, height, bits_per_comp, color_comp, pixel_data) = match format {
ImageFormat::Jpeg => {
let (w, h) = parse_jpeg_dimensions(&data)?;
(w, h, 8, 3, data)
}
ImageFormat::Png => parse_png_full(&data)?,
ImageFormat::Bmp => parse_bmp_full(&data)?,
};
Ok(ImageInfo {
format,
width,
height,
data: pixel_data,
bits_per_component: bits_per_comp,
color_components: color_comp,
alt_text,
})
}
impl ImageInfo {
pub fn with_alt_text(mut self, alt_text: String) -> Self {
self.alt_text = Some(alt_text);
self
}
pub fn get_alt_text(&self) -> &str {
self.alt_text.as_deref().unwrap_or("Image")
}
}
fn parse_png_full(data: &[u8]) -> Result<(u32, u32, u8, u8, Vec<u8>)> {
if data.len() < 24 {
return Err(anyhow!("PNG data too short"));
}
let width = u32::from_be_bytes([data[16], data[17], data[18], data[19]]);
let height = u32::from_be_bytes([data[20], data[21], data[22], data[23]]);
let bit_depth = data[24];
let color_type = data[25];
let (color_components, has_alpha) = match color_type {
0 => (1, false),
2 => (3, false),
3 => return Err(anyhow!("Paletted PNG (color type 3) not yet supported")),
4 => (2, true),
6 => (4, true),
_ => return Err(anyhow!("Invalid PNG color type: {}", color_type)),
};
let idat_data = extract_png_idat_chunks(data)?;
let decompressed = decompress_png_data(&idat_data)?;
let final_data = if has_alpha {
remove_alpha_channel(&decompressed, color_components, width, height)?
} else {
decompressed
};
Ok((width, height, bit_depth, color_components, final_data))
}
fn extract_png_idat_chunks(data: &[u8]) -> Result<Vec<u8>> {
let mut idat_data = Vec::new();
let mut i = 8;
while i + 8 <= data.len() {
let chunk_length = u32::from_be_bytes([data[i], data[i + 1], data[i + 2], data[i + 3]]) as usize;
let chunk_type = &data[i + 4..i + 8];
let chunk_data_start = i + 8;
let chunk_data_end = chunk_data_start + chunk_length;
if chunk_data_end > data.len() {
return Err(anyhow!("PNG chunk data extends beyond file"));
}
let chunk_type_str = std::str::from_utf8(chunk_type)
.map_err(|_| anyhow!("Invalid PNG chunk type"))?;
if chunk_type_str == "IDAT" {
idat_data.extend_from_slice(&data[chunk_data_start..chunk_data_end]);
} else if chunk_type_str == "IEND" {
break;
}
i = chunk_data_end + 4; }
if idat_data.is_empty() {
return Err(anyhow!("No IDAT chunks found in PNG"));
}
Ok(idat_data)
}
fn decompress_png_data(compressed: &[u8]) -> Result<Vec<u8>> {
crate::compression::decompress_deflate(compressed)
}
fn remove_alpha_channel(data: &[u8], components: u8, width: u32, height: u32) -> Result<Vec<u8>> {
let components = components as usize;
let bytes_per_pixel = components;
let _stride = width as usize * bytes_per_pixel + 1; let row_size = width as usize * components;
let mut result = Vec::new();
let mut i = 0;
for _ in 0..height {
if i + 1 > data.len() {
return Err(anyhow!("PNG data truncated"));
}
let filter = data[i];
i += 1;
if i + row_size > data.len() {
return Err(anyhow!("PNG row data truncated"));
}
result.push(filter);
let mut pixel_start = i;
for _ in 0..width as usize {
if pixel_start + components > data.len() {
return Err(anyhow!("PNG pixel data truncated"));
}
for c in 0..3 {
if c < components - 1 {
result.push(data[pixel_start + c]);
}
}
pixel_start += components;
}
i += row_size;
}
Ok(result)
}
fn parse_jpeg_dimensions(data: &[u8]) -> Result<(u32, u32)> {
let mut i = 2; while i + 1 < data.len() {
if data[i] != 0xFF {
i += 1;
continue;
}
let marker = data[i + 1];
i += 2;
if marker == 0xC0 || marker == 0xC1 || marker == 0xC2 {
if i + 7 > data.len() {
return Err(anyhow!("JPEG SOF marker truncated"));
}
let height = ((data[i + 3] as u32) << 8) | (data[i + 4] as u32);
let width = ((data[i + 5] as u32) << 8) | (data[i + 6] as u32);
return Ok((width, height));
}
if i + 1 >= data.len() {
break;
}
let seg_len = ((data[i] as usize) << 8) | (data[i + 1] as usize);
i += seg_len;
}
Err(anyhow!("Could not find JPEG SOF marker"))
}
fn parse_png_dimensions(data: &[u8]) -> Result<(u32, u32)> {
if data.len() < 24 {
return Err(anyhow!("PNG data too short"));
}
let width = u32::from_be_bytes([data[16], data[17], data[18], data[19]]);
let height = u32::from_be_bytes([data[20], data[21], data[22], data[23]]);
Ok((width, height))
}
fn parse_bmp_dimensions(data: &[u8]) -> Result<(u32, u32)> {
if data.len() < 26 {
return Err(anyhow!("BMP data too short"));
}
let width = u32::from_le_bytes([data[18], data[19], data[20], data[21]]);
let height_raw = i32::from_le_bytes([data[22], data[23], data[24], data[25]]);
let height = height_raw.unsigned_abs();
Ok((width, height))
}
fn parse_bmp_full(data: &[u8]) -> Result<(u32, u32, u8, u8, Vec<u8>)> {
if data.len() < 54 {
return Err(anyhow!("BMP data too short for header"));
}
let width = u32::from_le_bytes([data[18], data[19], data[20], data[21]]);
let height_raw = i32::from_le_bytes([data[22], data[23], data[24], data[25]]);
let height = height_raw.unsigned_abs();
let bits_per_pixel = u16::from_le_bytes([data[28], data[29]]);
let (bytes_per_pixel, _has_alpha) = match bits_per_pixel {
24 => (3, false),
32 => (4, true),
_ => return Err(anyhow!("Unsupported BMP bit depth: {} (only 24/32 supported)", bits_per_pixel)),
};
let row_size = ((width as usize * bytes_per_pixel + 3) / 4) * 4;
let pixel_data_offset = u32::from_le_bytes([data[10], data[11], data[12], data[13]]) as usize;
if pixel_data_offset as usize + row_size * height as usize > data.len() {
return Err(anyhow!("BMP pixel data truncated"));
}
let mut pixel_data = Vec::with_capacity((width * height * 3) as usize);
for y in (0..height as usize).rev() {
let row_start = pixel_data_offset + y * row_size;
for x in 0..width as usize {
let pixel_start = row_start + x * bytes_per_pixel;
let b = data[pixel_start];
let g = data[pixel_start + 1];
let r = data[pixel_start + 2];
pixel_data.push(r);
pixel_data.push(g);
pixel_data.push(b);
}
}
Ok((width, height, 8, 3, pixel_data))
}
pub fn scale_to_fit(width: u32, height: u32, max_width: f32, max_height: f32) -> (f32, f32) {
let w = width as f32;
let h = height as f32;
let scale_w = max_width / w;
let scale_h = max_height / h;
let scale = scale_w.min(scale_h).min(1.0); (w * scale, h * scale)
}
pub fn create_jpeg_image_object(
generator: &mut crate::pdf_generator::PdfGenerator,
jpeg_data: Vec<u8>,
width: u32,
height: u32,
) -> u32 {
let image_dict = format!(
"<< /Type /XObject\n\
/Subtype /Image\n\
/Width {}\n\
/Height {}\n\
/BitsPerComponent 8\n\
/ColorSpace /DeviceRGB\n\
/Filter /DCTDecode\n\
/Length {}\n\
>>\n",
width, height, jpeg_data.len()
);
generator.add_stream_object(image_dict, jpeg_data)
}
pub fn create_png_image_object(
generator: &mut crate::pdf_generator::PdfGenerator,
png_data: Vec<u8>,
width: u32,
height: u32,
bits_per_component: u8,
color_components: u8,
) -> u32 {
let color_space = match color_components {
1 => "/DeviceGray",
3 => "/DeviceRGB",
_ => "/DeviceRGB", };
let image_dict = format!(
"<< /Type /XObject\n\
/Subtype /Image\n\
/Width {}\n\
/Height {}\n\
/BitsPerComponent {}\n\
/ColorSpace {}\n\
/Filter /FlateDecode\n\
/DecodeParms << /Predictor 15 /Colors {} /BitsPerComponent {} /Columns {} >>\n\
/Length {}\n\
>>\n",
width, height, bits_per_component, color_space,
color_components, bits_per_component, width, png_data.len()
);
generator.add_stream_object(image_dict, png_data)
}
pub fn create_bmp_image_object(
generator: &mut crate::pdf_generator::PdfGenerator,
bmp_data: Vec<u8>,
width: u32,
height: u32,
) -> u32 {
let image_dict = format!(
"<< /Type /XObject\n\
/Subtype /Image\n\
/Width {}\n\
/Height {}\n\
/BitsPerComponent 8\n\
/ColorSpace /DeviceRGB\n\
/Length {}\n\
>>\n",
width, height, bmp_data.len()
);
generator.add_stream_object(image_dict, bmp_data)
}
pub fn create_image_object(
generator: &mut crate::pdf_generator::PdfGenerator,
image_info: ImageInfo,
) -> Result<u32> {
match image_info.format {
ImageFormat::Jpeg => {
Ok(create_jpeg_image_object(
generator,
image_info.data,
image_info.width,
image_info.height,
))
}
ImageFormat::Png => {
Ok(create_png_image_object(
generator,
image_info.data,
image_info.width,
image_info.height,
image_info.bits_per_component,
image_info.color_components,
))
}
ImageFormat::Bmp => {
Ok(create_bmp_image_object(
generator,
image_info.data,
image_info.width,
image_info.height,
))
}
}
}
pub fn create_image_content_stream(
x: f32,
y: f32,
width: f32,
height: f32,
image_name: &str,
) -> Vec<u8> {
let mut content = Vec::new();
content.extend_from_slice(b"q\n");
content.extend_from_slice(
format!("{} 0 0 {} {} {} cm\n", width, height, x, y).as_bytes(),
);
content.extend_from_slice(format!("/{} Do\n", image_name).as_bytes());
content.extend_from_slice(b"Q\n");
content
}
pub fn add_image_to_pdf(
output_pdf: &str,
image_path: &str,
x: f32,
y: f32,
display_width: f32,
display_height: f32,
) -> Result<()> {
let info = load_image(image_path)?;
let mut generator = crate::pdf_generator::PdfGenerator::new();
let image_id = create_image_object(&mut generator, info.clone())?;
let content = create_image_content_stream(x, y, display_width, display_height, "Im1");
let content_id = generator.add_stream_object(
format!("<< /Length {} >>\n", content.len()),
content,
);
let page_dict = format!(
"<< /Type /Page\n\
/Parent 5 0 R\n\
/MediaBox [0 0 612 792]\n\
/Contents {} 0 R\n\
/Resources << /XObject << /Im1 {} 0 R >> >>\n\
>>\n",
content_id, image_id
);
let page_id = generator.add_object(page_dict);
let pages_dict = format!(
"<< /Type /Pages\n/Kids [{} 0 R]\n/Count 1\n>>\n",
page_id
);
let pages_id = generator.add_object(pages_dict);
let catalog = format!("<< /Type /Catalog\n/Pages {} 0 R\n>>\n", pages_id);
generator.add_object(catalog);
let pdf_data = generator.generate();
std::fs::write(output_pdf, &pdf_data)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_detect_jpeg() {
let data = vec![0xFF, 0xD8, 0xFF, 0xE0, 0x00];
assert_eq!(detect_image_format(&data).unwrap(), ImageFormat::Jpeg);
}
#[test]
fn test_detect_png() {
let data = vec![0x89, 0x50, 0x4E, 0x47, 0x0D];
assert_eq!(detect_image_format(&data).unwrap(), ImageFormat::Png);
}
#[test]
fn test_detect_bmp() {
let data = vec![0x42, 0x4D, 0x00, 0x00];
assert_eq!(detect_image_format(&data).unwrap(), ImageFormat::Bmp);
}
#[test]
fn test_detect_unknown() {
let data = vec![0x00, 0x00, 0x00, 0x00];
assert!(detect_image_format(&data).is_err());
}
#[test]
fn test_scale_to_fit() {
let (w, h) = scale_to_fit(800, 600, 400.0, 400.0);
assert!((w - 400.0).abs() < 0.01);
assert!((h - 300.0).abs() < 0.01);
}
#[test]
fn test_scale_no_upscale() {
let (w, h) = scale_to_fit(100, 50, 400.0, 400.0);
assert!((w - 100.0).abs() < 0.01);
assert!((h - 50.0).abs() < 0.01);
}
#[test]
fn test_parse_jpeg_dimensions() {
let mut data = vec![0xFF, 0xD8]; data.extend_from_slice(&[0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00]);
data.extend_from_slice(&[0xFF, 0xC0]);
data.extend_from_slice(&[0x00, 0x11]); data.push(0x08); data.extend_from_slice(&[0x01, 0x00]); data.extend_from_slice(&[0x02, 0x00]); data.extend_from_slice(&[0x03]); data.extend_from_slice(&[0; 20]);
let (w, h) = parse_jpeg_dimensions(&data).unwrap();
assert_eq!(w, 512);
assert_eq!(h, 256);
}
#[test]
fn test_parse_png_dimensions() {
let mut data = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; data.extend_from_slice(&[0x00, 0x00, 0x00, 0x0D]); data.extend_from_slice(b"IHDR");
data.extend_from_slice(&640u32.to_be_bytes()); data.extend_from_slice(&480u32.to_be_bytes());
let (w, h) = parse_png_dimensions(&data).unwrap();
assert_eq!(w, 640);
assert_eq!(h, 480);
}
#[test]
fn test_create_image_content_stream() {
let cs = create_image_content_stream(100.0, 200.0, 300.0, 400.0, "Im1");
let s = String::from_utf8(cs).unwrap();
assert!(s.contains("q\n"));
assert!(s.contains("300 0 0 400 100 200 cm"));
assert!(s.contains("/Im1 Do"));
assert!(s.contains("Q\n"));
}
#[test]
fn test_png_color_components() {
assert_eq!(get_png_color_components(0), Some((1, false)));
assert_eq!(get_png_color_components(2), Some((3, false)));
assert_eq!(get_png_color_components(4), Some((2, true)));
assert_eq!(get_png_color_components(6), Some((4, true)));
assert_eq!(get_png_color_components(1), None);
assert_eq!(get_png_color_components(5), None);
}
#[test]
fn test_bmp_bit_depth_validation() {
assert!(validate_bmp_bit_depth(24).is_ok());
assert!(validate_bmp_bit_depth(32).is_ok());
assert!(validate_bmp_bit_depth(8).is_err());
assert!(validate_bmp_bit_depth(16).is_err());
assert!(validate_bmp_bit_depth(1).is_err());
}
#[test]
fn test_bmp_row_padding() {
let row_size = calculate_bmp_row_size(1, 3);
assert_eq!(row_size, 4);
let row_size = calculate_bmp_row_size(2, 3);
assert_eq!(row_size, 8);
let row_size = calculate_bmp_row_size(3, 3);
assert_eq!(row_size, 12);
let row_size = calculate_bmp_row_size(4, 3);
assert_eq!(row_size, 12);
}
}
fn get_png_color_components(color_type: u8) -> Option<(u8, bool)> {
match color_type {
0 => Some((1, false)), 2 => Some((3, false)), 4 => Some((2, true)), 6 => Some((4, true)), _ => None,
}
}
fn validate_bmp_bit_depth(bits_per_pixel: u16) -> Result<()> {
match bits_per_pixel {
24 | 32 => Ok(()),
_ => Err(anyhow!("Unsupported BMP bit depth: {}", bits_per_pixel)),
}
}
fn calculate_bmp_row_size(width: u32, bytes_per_pixel: u8) -> usize {
let row_size = width as usize * bytes_per_pixel as usize;
((row_size + 3) / 4) * 4
}
#[cfg(test)]
mod proptest_tests {
use super::*;
use proptest::prelude::*;
proptest! {
#[test]
fn scale_preserves_aspect_ratio(width in 1u32..4000u32, height in 1u32..4000u32,
max_w in 100f32..2000f32, max_h in 100f32..2000f32) {
let (scaled_w, scaled_h) = scale_to_fit(width, height, max_w, max_h);
assert!(scaled_w <= max_w + 0.01, "Scaled width exceeds max");
assert!(scaled_h <= max_h + 0.01, "Scaled height exceeds max");
let original_aspect = width as f32 / height as f32;
let scaled_aspect = scaled_w / scaled_h;
assert!((original_aspect - scaled_aspect).abs() < 0.01f32, "Aspect ratio not preserved");
assert!(scaled_w <= width as f32 + 0.01, "Width was upscaled");
assert!(scaled_h <= height as f32 + 0.01, "Height was upscaled");
}
}
proptest! {
#[test]
fn scale_never_exceeds_bounds(width in 1u32..4000u32, height in 1u32..4000u32,
max_w in 100f32..2000f32, max_h in 100f32..2000f32) {
let (scaled_w, scaled_h) = scale_to_fit(width, height, max_w, max_h);
assert!(scaled_w <= max_w + 0.01, "Scaled width {} exceeds max_w {}", scaled_w, max_w);
assert!(scaled_h <= max_h + 0.01, "Scaled height {} exceeds max_h {}", scaled_h, max_h);
}
}
}