playa 0.1.142

//! Image loader with pluggable backends
//!
//! Unified interface for loading image files with metadata extraction.
//! Supports different backends based on feature flags:
//! - Default: `image` crate (uses exrs for EXR)
//! - Feature "openexr": openexr-rs (C++ bindings, full DWAA/DWAB support)

#[cfg_attr(feature = "openexr", allow(unused_imports))]
use half::f16 as F16;
use log::trace;
use std::path::Path;

use super::frame::{Frame, FrameError, PixelBuffer, PixelFormat};
use crate::entities::loader_video;
use crate::entities::{AttrValue, Attrs};
use crate::utils::media;
use super::keys::{A_FPS, A_HEIGHT, A_WIDTH};

/// Image loader with metadata support
pub struct Loader;

impl Loader {
    /// Read image file header and extract metadata
    ///
    /// Returns Attrs with metadata like:
    /// - "width", "height" (UInt) - image dimensions
    /// - "channels" (UInt) - number of channels
    /// - "format" (Str) - pixel format description
    /// - Additional format-specific metadata
    pub fn header(path: &Path) -> Result<Attrs, FrameError> {
        let ext = path
            .extension()
            .and_then(|s| s.to_str())
            .unwrap_or("")
            .to_lowercase();

        match ext.as_str() {
            ext if media::VIDEO_EXTS.contains(&ext) => Self::header_video(path),
            "exr" => Self::header_exr(path),
            _ => Self::header_generic(path),
        }
    }

    /// Load complete image file into Frame
    pub fn load(path: &Path) -> Result<Frame, FrameError> {
        let ext = path
            .extension()
            .and_then(|s| s.to_str())
            .unwrap_or("")
            .to_lowercase();

        match ext.as_str() {
            ext if media::VIDEO_EXTS.contains(&ext) => Self::load_video(path),
            "exr" => Self::load_exr(path),
            _ => Self::load_generic(path),
        }
    }

    // ===== Video Loading =====

    /// Read video metadata into Attrs (width, height, fps, frames)
    fn header_video(path: &Path) -> Result<Attrs, FrameError> {
        let (actual_path, _) = media::parse_video_path(path);
        let meta = loader_video::VideoMetadata::from_file(&actual_path)?;

        let mut meta_attrs = Attrs::new();
        meta_attrs.set(A_WIDTH, AttrValue::UInt(meta.width));
        meta_attrs.set(A_HEIGHT, AttrValue::UInt(meta.height));
        meta_attrs.set(
            "format",
            AttrValue::Str(format!("Video ({})", actual_path.display())),
        );
        meta_attrs.set("channels", AttrValue::UInt(3));
        meta_attrs.set("frames", AttrValue::UInt(meta.frame_count as u32));
        meta_attrs.set(A_FPS, AttrValue::Float(meta.fps as f32));

        Ok(meta_attrs)
    }

    /// Load a single video frame into Frame (defaults to frame 0 if not specified)
    fn load_video(path: &Path) -> Result<Frame, FrameError> {
        let (actual_path, frame_idx) = media::parse_video_path(path);
        let frame_num = frame_idx.unwrap_or(0);
        let (buffer, pixel_format, width, height) =
            loader_video::decode_frame(&actual_path, frame_num)?;

        Ok(Frame::from_buffer(buffer, pixel_format, width, height))
    }

    // ===== EXR Loading =====

    /// EXR header reading with openexr-rs (C++ bindings)
    #[cfg(feature = "openexr")]
    fn header_exr(path: &Path) -> Result<Attrs, FrameError> {
        trace!("Reading EXR header with openexr-rs: {}", path.display());
        use openexr::prelude::*;

        let file = RgbaInputFile::new(path, 1)
            .map_err(|e| FrameError::Image(format!("OpenEXR header error: {}", e)))?;

        let header = file.header();
        let data_window = header.data_window::<[i32; 4]>();
        let width = (data_window[2] - data_window[0] + 1) as usize;
        let height = (data_window[3] - data_window[1] + 1) as usize;

        // Count channels via iterator
        let channels = header.channels();
        let channel_count = channels.iter().count();

        let mut meta = Attrs::new();
        meta.set(A_WIDTH, AttrValue::UInt(width as u32));
        meta.set(A_HEIGHT, AttrValue::UInt(height as u32));
        meta.set("format", AttrValue::Str("EXR (OpenEXR)".to_string()));
        meta.set("channels", AttrValue::UInt(channel_count as u32));

        Ok(meta)
    }

    #[cfg(not(feature = "openexr"))]
    fn header_exr(path: &Path) -> Result<Attrs, FrameError> {
        trace!("Reading EXR header with image crate: {}", path.display());

        // Use image crate for header reading (it uses exrs internally)
        let reader = image::ImageReader::open(path)
            .map_err(|e| FrameError::Image(format!("Failed to open EXR: {}", e)))?;

        let format = reader
            .format()
            .ok_or_else(|| FrameError::Image("Failed to detect image format".to_string()))?;

        let img = reader.decode().map_err(|e| {
            let err_str = e.to_string();
            if err_str.contains("DWAA") || err_str.contains("DWAB") {
                FrameError::UnsupportedFormat(
                    "DWAA/DWAB compression not supported. Build with: cargo xtask build --openexr"
                        .to_string(),
                )
            } else {
                FrameError::Image(format!("EXR decode error: {}", e))
            }
        })?;

        let mut meta = Attrs::new();
        meta.set(A_WIDTH, AttrValue::UInt(img.width()));
        meta.set(A_HEIGHT, AttrValue::UInt(img.height()));
        meta.set("format", AttrValue::Str(format!("EXR ({:?})", format)));

        // Determine channel count from color type
        let channels = match img.color() {
            image::ColorType::L8 | image::ColorType::L16 => 1,
            image::ColorType::La8 | image::ColorType::La16 => 2,
            image::ColorType::Rgb8 | image::ColorType::Rgb16 | image::ColorType::Rgb32F => 3,
            image::ColorType::Rgba8 | image::ColorType::Rgba16 | image::ColorType::Rgba32F => 4,
            _ => 4,
        };
        meta.set("channels", AttrValue::UInt(channels));

        Ok(meta)
    }

    /// EXR loading with openexr-rs - detects pixel type and uses optimal path
    #[cfg(feature = "openexr")]
    fn load_exr(path: &Path) -> Result<Frame, FrameError> {
        trace!("Loading EXR with openexr-rs: {}", path.display());
        use openexr::prelude::*;

        // Open file to detect pixel type from R channel
        let file = RgbaInputFile::new(path, 1)
            .map_err(|e| FrameError::Image(format!("OpenEXR error: {}", e)))?;

        let header = file.header();
        let data_window = header.data_window::<[i32; 4]>();
        let width = (data_window[2] - data_window[0] + 1) as usize;
        let height = (data_window[3] - data_window[1] + 1) as usize;

        // Detect pixel type from R channel (HALF, FLOAT, or UINT)
        let channels = header.channels();
        let pixel_type = channels
            .iter()
            .find(|(name, _)| *name == "R")
            .map(|(_, ch)| ch.type_)
            .unwrap_or(PixelType::Half.into());

        drop(header);
        drop(file);

        // Dispatch to optimal loader based on pixel type
        if pixel_type == PixelType::Float.into() {
            Self::load_exr_float(path, width, height)
        } else {
            // HALF or UINT - load as f16 (native for HALF, memory-efficient for UINT)
            Self::load_exr_half(path, width, height)
        }
    }

    /// Load EXR with HALF pixels using RgbaInputFile (native f16)
    /// Converts openexr's half 1.x to our half 2.x via raw bits
    #[cfg(feature = "openexr")]
    fn load_exr_half(path: &Path, width: usize, height: usize) -> Result<Frame, FrameError> {
        use openexr::prelude::*;

        let mut file = RgbaInputFile::new(path, 1)
            .map_err(|e| FrameError::Image(format!("OpenEXR error: {}", e)))?;

        let header = file.header();
        let data_window = header.data_window::<[i32; 4]>();
        let y_min = data_window[1];
        let y_max = data_window[3];
        drop(header);

        // Read as Rgba (uses half 1.x f16 internally)
        let mut pixels_rgba = vec![Rgba::from_f32(0.0, 0.0, 0.0, 0.0); width * height];
        file.set_frame_buffer(&mut pixels_rgba, 1, width)
            .map_err(|e| FrameError::Image(format!("OpenEXR framebuffer error: {}", e)))?;

        unsafe {
            file.read_pixels(y_min, y_max)
                .map_err(|e| FrameError::Image(format!("OpenEXR read error: {}", e)))?;
        }

        // Convert half 1.x -> half 2.x via raw bits (same binary format)
        let pixel_count = width * height;
        let mut buffer: Vec<half::f16> = Vec::with_capacity(pixel_count * 4);

        for pixel in pixels_rgba.iter() {
            // Extract raw u16 bits from half 1.x, create half 2.x
            buffer.push(half::f16::from_bits(pixel.r.to_bits()));
            buffer.push(half::f16::from_bits(pixel.g.to_bits()));
            buffer.push(half::f16::from_bits(pixel.b.to_bits()));
            buffer.push(half::f16::from_bits(pixel.a.to_bits()));
        }

        trace!("Loaded EXR HALF: {}x{} (f16)", width, height);
        Ok(super::frame::Frame::from_buffer(
            PixelBuffer::F16(buffer),
            PixelFormat::RgbaF16,
            width,
            height,
        ))
    }

    /// Load EXR with FLOAT pixels using Frame API (native f32, full precision)
    #[cfg(feature = "openexr")]
    fn load_exr_float(path: &Path, width: usize, height: usize) -> Result<Frame, FrameError> {
        use openexr::prelude::*;

        let file = InputFile::new(path, 1)
            .map_err(|e| FrameError::Image(format!("OpenEXR error: {}", e)))?;

        let header = file.header();
        let data_window = *header.data_window::<[i32; 4]>();
        let y_min = data_window[1];
        let y_max = data_window[3];
        drop(header);

        // Create Frame with f32 for RGBA channels (native precision, no f16 conversion)
        let frame_rgba = Frame::new::<f32, _, _>(&["R", "G", "B", "A"], data_window)
            .map_err(|e| FrameError::Image(format!("OpenEXR frame error: {}", e)))?;

        // Read pixels into frame
        let (_file, mut frames) = file
            .into_reader(vec![frame_rgba])
            .map_err(|e| FrameError::Image(format!("OpenEXR reader error: {}", e)))?
            .read_pixels(y_min, y_max)
            .map_err(|e| FrameError::Image(format!("OpenEXR read error: {}", e)))?;

        // Extract flat RGBA f32 buffer
        let buffer_f32: Vec<f32> = frames.remove(0).into_vec();

        trace!("Loaded EXR FLOAT: {}x{} (f32, native precision)", width, height);
        Ok(super::frame::Frame::from_buffer(
            PixelBuffer::F32(buffer_f32),
            PixelFormat::RgbaF32,
            width,
            height,
        ))
    }

    #[cfg(not(feature = "openexr"))]
    fn load_exr(path: &Path) -> Result<Frame, FrameError> {
        trace!("Loading EXR with image crate: {}", path.display());

        let img = image::open(path).map_err(|e| {
            let err_str = e.to_string();
            if err_str.contains("DWAA") || err_str.contains("DWAB") {
                return FrameError::UnsupportedFormat(
                    "DWAA/DWAB compression not supported. Build with: cargo xtask build --openexr"
                        .to_string(),
                );
            }
            FrameError::Image(format!("EXR load error: {}", e))
        })?;

        let width = img.width() as usize;
        let height = img.height() as usize;

        // Convert to Rgba32F
        let rgba_img = img.to_rgba32f();
        let pixels = rgba_img.as_raw();

        // Convert f32 to f16
        let mut buffer = Vec::with_capacity(pixels.len());
        for &pixel in pixels {
            buffer.push(F16::from_f32(pixel));
        }

        Ok(Frame::from_buffer(
            PixelBuffer::F16(buffer),
            PixelFormat::RgbaF16,
            width,
            height,
        ))
    }

    // ===== Generic Image Loading (PNG, JPEG, TIFF, etc.) =====

    fn header_generic(path: &Path) -> Result<Attrs, FrameError> {
        trace!("Reading generic image header: {}", path.display());

        let reader = image::ImageReader::open(path)
            .map_err(|e| FrameError::Image(format!("Failed to open image: {}", e)))?;

        let format = reader
            .format()
            .ok_or_else(|| FrameError::Image("Failed to detect image format".to_string()))?;

        let img = reader
            .decode()
            .map_err(|e| FrameError::Image(format!("Image decode error: {}", e)))?;

        let mut meta = Attrs::new();
        meta.set(A_WIDTH, AttrValue::UInt(img.width()));
        meta.set(A_HEIGHT, AttrValue::UInt(img.height()));
        meta.set("format", AttrValue::Str(format!("{:?}", format)));

        let channels = match img.color() {
            image::ColorType::L8 | image::ColorType::L16 => 1,
            image::ColorType::La8 | image::ColorType::La16 => 2,
            image::ColorType::Rgb8 | image::ColorType::Rgb16 | image::ColorType::Rgb32F => 3,
            image::ColorType::Rgba8 | image::ColorType::Rgba16 | image::ColorType::Rgba32F => 4,
            _ => 4,
        };
        meta.set("channels", AttrValue::UInt(channels));

        Ok(meta)
    }

    fn load_generic(path: &Path) -> Result<Frame, FrameError> {
        trace!("Loading generic image: {}", path.display());

        let img =
            image::open(path).map_err(|e| FrameError::Image(format!("Image load error: {}", e)))?;

        let width = img.width() as usize;
        let height = img.height() as usize;

        // Convert to Rgba8
        let rgba_img = img.to_rgba8();
        let pixels = rgba_img.into_raw();

        Ok(Frame::from_buffer(
            PixelBuffer::U8(pixels),
            PixelFormat::Rgba8,
            width,
            height,
        ))
    }
}