pbrt/src/core/image/io.rs

use super::{Image, ImageAndMetadata, ImageMetadata};
use crate::core::image::{PixelStorage, WrapMode};
use crate::utils::error::ImageError;
use anyhow::Error;
use anyhow::{Context, Result, bail};
use exr::prelude::{read_first_rgba_layer_from_file, write_rgba_file};
use image_rs::{DynamicImage, ImageReader};
use shared::Float;
use shared::core::color::{ColorEncoding, LINEAR};
use shared::core::geometry::Point2i;
use shared::core::image::{DeviceImage, ImageBase, PixelFormat};
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Write};
use std::path::Path;

pub trait ImageIO {
    fn read(path: &Path, encoding: Option<ColorEncoding>) -> Result<ImageAndMetadata>;
    fn write(&self, filename: &str, metadata: &ImageMetadata) -> Result<()>;
    fn write_png(&self, path: &Path) -> Result<()>;
    fn write_exr(&self, path: &Path, metadata: &ImageMetadata) -> Result<()>;
    fn write_qoi(&self, path: &Path) -> Result<()>;
    fn write_pfm(&self, path: &Path) -> Result<()>;
    fn to_u8_buffer(&self) -> Vec<u8>;
}

impl ImageIO for Image {
    fn read(path: &Path, encoding: Option<ColorEncoding>) -> Result<ImageAndMetadata> {
        let ext = path
            .extension()
            .and_then(|s| s.to_str())
            .unwrap_or("")
            .to_lowercase();

        match ext.as_str() {
            "exr" => read_exr(path),
            "pfm" => read_pfm(path),
            _ => read_generic(path, encoding),
        }
    }

    fn write(&self, filename: &str, metadata: &ImageMetadata) -> Result<(), Error> {
        let path = Path::new(filename);
        let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
        let res = match ext.to_lowercase().as_str() {
            "exr" => self.write_exr(path, metadata),
            "png" => self.write_png(path),
            "pfm" => self.write_pfm(path),
            "qoi" => self.write_qoi(path),
            _ => Err(anyhow::anyhow!("Unsupported write format: {}", ext)),
        };
        res.map_err(|e| ImageError::Io(std::io::Error::other(e)))
    }

    fn write_png(&self, path: &Path) -> Result<()> {
        let w = self.resolution.x() as u32;
        let h = self.resolution.y() as u32;

        // Convert whatever we have to u8 [0..255]
        let data = self.to_u8_buffer();
        let channels = self.n_channels();

        match channels {
            1 => {
                // Luma
                image_rs::save_buffer_with_format(
                    path,
                    &data,
                    w,
                    h,
                    image_rs::ColorType::L8,
                    image_rs::ImageFormat::Png,
                )?;
            }
            3 => {
                // RGB
                image_rs::save_buffer_with_format(
                    path,
                    &data,
                    w,
                    h,
                    image_rs::ColorType::Rgb8,
                    image_rs::ImageFormat::Png,
                )?;
            }
            4 => {
                // RGBA
                image_rs::save_buffer_with_format(
                    path,
                    &data,
                    w,
                    h,
                    image_rs::ColorType::Rgba8,
                    image_rs::ImageFormat::Png,
                )?;
            }
            _ => bail!("PNG writer only supports 1, 3, or 4 channels"),
        }
        Ok(())
    }

    fn write_qoi(&self, path: &Path) -> Result<()> {
        let w = self.resolution.x() as u32;
        let h = self.resolution.y() as u32;
        let data = self.to_u8_buffer();

        let color_type = match self.n_channels() {
            3 => image_rs::ColorType::Rgb8,
            4 => image_rs::ColorType::Rgba8,
            _ => bail!("QOI only supports 3 or 4 channels"),
        };

        image_rs::save_buffer_with_format(
            path,
            &data,
            w,
            h,
            color_type,
            image_rs::ImageFormat::Qoi,
        )?;
        Ok(())
    }

    fn write_exr(&self, path: &Path, _metadata: &ImageMetadata) -> Result<()> {
        // EXR requires F32
        let w = self.resolution.x() as usize;
        let h = self.resolution.y() as usize;
        let c = self.n_channels();

        write_rgba_file(path, w, h, |x, y| {
            // Helper to get float value regardless of internal storage
            let get = |ch| {
                self.get_channel_with_wrap(
                    Point2i::new(x as i32, y as i32),
                    ch,
                    WrapMode::Clamp.into(),
                )
            };

            if c == 1 {
                let v = get(0);
                (v, v, v, 1.0)
            } else if c == 3 {
                (get(0), get(1), get(2), 1.0)
            } else {
                (get(0), get(1), get(2), get(3))
            }
        })
        .context("Failed to write EXR")?;

        Ok(())
    }

    fn write_pfm(&self, path: &Path) -> Result<()> {
        let file = File::create(path)?;
        let mut writer = BufWriter::new(file);

        if self.n_channels() != 3 {
            bail!("PFM writing currently only supports 3 channels (RGB)");
        }

        // Header
        writeln!(writer, "PF")?;
        writeln!(writer, "{} {}", self.resolution.x(), self.resolution.y())?;
        let scale = if cfg!(target_endian = "little") {
            -1.0
        } else {
            1.0
        };
        writeln!(writer, "{}", scale)?;

        // PBRT stores top-to-bottom.
        for y in (0..self.resolution.y()).rev() {
            for x in 0..self.resolution.x() {
                for c in 0..3 {
                    let val =
                        self.get_channel_with_wrap(Point2i::new(x, y), c, WrapMode::Clamp.into());
                    writer.write_all(&val.to_le_bytes())?;
                }
            }
        }

        Ok(())
    }

    fn to_u8_buffer(&self) -> Vec<u8> {
        match &self.pixels {
            PixelStorage::U8(data) => data.clone(),
            PixelStorage::F16(data) => data
                .iter()
                .map(|v| (v.to_f32().clamp(0.0, 1.0) * 255.0 + 0.5) as u8)
                .collect(),
            PixelStorage::F32(data) => data
                .iter()
                .map(|v| (v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8)
                .collect(),
        }
    }
}

fn read_generic(path: &Path, encoding: Option<ColorEncoding>) -> Result<ImageAndMetadata> {
    let dyn_img = ImageReader::open(path)
        .with_context(|| format!("Failed to open image: {:?}", path))?
        .decode()?;

    let w = dyn_img.width() as i32;
    let h = dyn_img.height() as i32;
    let res = Point2i::new(w, h);

    // Check if it was loaded as high precision or standard
    let rgb_names = || vec!["R".to_string(), "G".to_string(), "B".to_string()];
    let rgba_names = || {
        vec![
            "R".to_string(),
            "G".to_string(),
            "B".to_string(),
            "A".to_string(),
        ]
    };
    let image = match dyn_img {
        DynamicImage::ImageRgb32F(buf) => Image::from_f32(buf.into_raw(), res, rgb_names()),
        DynamicImage::ImageRgba32F(buf) => Image::from_f32(buf.into_raw(), res, rgba_names()),
        _ => {
            // Default to RGB8 for everything else
            let enc = encoding.unwrap_or(ColorEncoding::sRGB);
            if dyn_img.color().has_alpha() {
                let buf = dyn_img.to_rgba8();
                Image::from_u8(buf.into_raw(), res, rgba_names(), enc)
            } else {
                let buf = dyn_img.to_rgb8();
                Image::from_u8(buf.into_raw(), res, rgb_names(), enc)
            }
        }
    };

    let metadata = ImageMetadata::default();
    Ok(ImageAndMetadata { image, metadata })
}

fn read_exr(path: &Path) -> Result<ImageAndMetadata> {
    let image = read_first_rgba_layer_from_file(
        path,
        |resolution, _| {
            let size = resolution.width() * resolution.height() * 4;
            vec![0.0 as Float; size]
        },
        |buffer, position, pixel| {
            let width = position.width();
            let idx = (position.y() * width + position.x()) * 4;
            // Map exr pixel struct to our buffer
            buffer[idx] = pixel.0;
            buffer[idx + 1] = pixel.1;
            buffer[idx + 2] = pixel.2;
            buffer[idx + 3] = pixel.3;
        },
    )
    .with_context(|| format!("Failed to read EXR: {:?}", path))?;

    let w = image.layer_data.size.width() as i32;
    let h = image.layer_data.size.height() as i32;

    let image = Image {
        format: PixelFormat::F32,
        resolution: Point2i::new(w, h),
        channel_names: vec!["R".into(), "G".into(), "B".into(), "A".into()],
        encoding: LINEAR,
        pixels: PixelStorage::F32(image.layer_data.channel_data.pixels),
    };

    let metadata = ImageMetadata::default();
    Ok(ImageAndMetadata { image, metadata })
}

fn read_pfm(path: &Path) -> Result<ImageAndMetadata> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);

    // PFM Headers are: "PF\nwidth height\nscale\n" (or Pf for grayscale)
    let mut header_word = String::new();
    reader.read_line(&mut header_word)?;
    let header_word = header_word.trim();

    let channels = match header_word {
        "PF" => 3,
        "Pf" => 1,
        _ => bail!("Invalid PFM header: {}", header_word),
    };

    let mut dims_line = String::new();
    reader.read_line(&mut dims_line)?;
    let dims: Vec<i32> = dims_line
        .split_whitespace()
        .map(|s| s.parse().unwrap_or(0))
        .collect();

    if dims.len() < 2 {
        bail!("Invalid PFM dimensions");
    }
    let w = dims[0];
    let h = dims[1];

    let mut scale_line = String::new();
    reader.read_line(&mut scale_line)?;
    let scale: f32 = scale_line.trim().parse().context("Invalid PFM scale")?;

    let file_is_little_endian = scale < 0.0;
    let abs_scale = scale.abs();

    let mut buffer = Vec::new();
    reader.read_to_end(&mut buffer)?;

    let expected_bytes = (w * h * channels) as usize * 4;
    if buffer.len() < expected_bytes {
        bail!("PFM file too short");
    }

    let mut pixels = vec![0.0 as Float; (w * h * channels) as usize];

    // PFM is Bottom-to-Top
    for y in 0..h {
        // Flippety-do
        let src_y = h - 1 - y;
        for x in 0..w {
            for c in 0..channels {
                let src_idx = ((src_y * w + x) * channels + c) as usize * 4;
                let dst_idx = ((y * w + x) * channels + c) as usize;

                let bytes: [u8; 4] = buffer[src_idx..src_idx + 4].try_into()?;

                let val = if file_is_little_endian {
                    f32::from_le_bytes(bytes)
                } else {
                    f32::from_be_bytes(bytes)
                };

                pixels[dst_idx] = val * abs_scale;
            }
        }
    }

    let names = if channels == 1 {
        vec!["Y".into()]
    } else {
        vec!["R".into(), "G".into(), "B".into()]
    };

    let image = Image::new(PixelFormat::F32, Point2i::new(w, h), names, LINEAR);
    let metadata = ImageMetadata::default();

    Ok(ImageAndMetadata { image, metadata })
}