zenjxl-decoder 0.3.8

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//! Official JPEG XL conformance tests.
//!
//! These tests compare jxl-rs output against the official libjxl conformance
//! test suite reference images (NPY format).
//!
//! To run these tests:
//! ```sh
//! cargo test --features cms conformance -- --ignored --nocapture
//! ```
//!
//! The conformance repo will be automatically cloned to `target/conformance/`
//! and reference data downloaded on first run. This may take a few minutes.
//!
//! To use a custom location, set `CONFORMANCE_PATH`:
//! ```sh
//! CONFORMANCE_PATH=/path/to/conformance cargo test --features cms conformance -- --ignored
//! ```
//!
//! IMPORTANT: DO NOT WEAKEN TOLERANCES. If a test fails, the implementation
//! is wrong and must be fixed.
#![allow(dead_code, unused_imports)] // Conformance infrastructure — used only with `cms` feature + --ignored tests

#[cfg(feature = "cms")]
use crate::api::MoxCms;
#[cfg(feature = "cms")]
use crate::api::{JxlCms, JxlColorProfile};
use crate::api::{
    JxlColorType, JxlDataFormat, JxlDecoder, JxlDecoderOptions, JxlOutputBuffer, JxlPixelFormat,
    ProcessingResult, states,
};

use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::Once;

/// One-time initialization for auto-fetching conformance data.
static INIT: Once = Once::new();

/// Get the conformance test directory.
/// Priority: CONFORMANCE_PATH env var > target/conformance
fn conformance_dir() -> Option<PathBuf> {
    // Check env var first
    if let Ok(path) = std::env::var("CONFORMANCE_PATH") {
        return Some(PathBuf::from(path));
    }

    // Use target/conformance as default (gitignored)
    let target_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()? // jxl-rs root
        .join("target")
        .join("conformance");

    Some(target_dir)
}

/// Ensure the conformance repo is cloned and reference data downloaded.
/// This is called once per test run.
fn ensure_conformance_data() -> Result<PathBuf, String> {
    let dir = conformance_dir().ok_or("Could not determine conformance directory")?;

    INIT.call_once(|| {
        if let Err(e) = setup_conformance_repo(&dir) {
            eprintln!("Warning: Failed to setup conformance repo: {}", e);
        }
    });

    if !dir.join("testcases").exists() {
        return Err(format!(
            "Conformance testcases not found at {:?}. Setup may have failed.",
            dir
        ));
    }

    Ok(dir)
}

/// Clone the conformance repo and download reference data if needed.
fn setup_conformance_repo(dir: &Path) -> Result<(), String> {
    let testcases_dir = dir.join("testcases");

    // Check if already set up
    if testcases_dir.exists() {
        // Check if reference data is downloaded (look for any .npy symlinks)
        let has_npy = std::fs::read_dir(&testcases_dir)
            .ok()
            .map(|entries| {
                entries.filter_map(|e| e.ok()).any(|e| {
                    std::fs::read_dir(e.path())
                        .ok()
                        .map(|inner| {
                            inner
                                .filter_map(|f| f.ok())
                                .any(|f| f.path().extension().is_some_and(|ext| ext == "npy"))
                        })
                        .unwrap_or(false)
                })
            })
            .unwrap_or(false);

        if has_npy {
            return Ok(()); // Already fully set up
        }

        // Need to download reference data
        eprintln!("Downloading conformance reference data...");
        return run_download_script(dir);
    }

    // Clone the repo
    eprintln!("Cloning libjxl/conformance to {:?}...", dir);

    // Create parent directory
    if let Some(parent) = dir.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("Failed to create directory: {}", e))?;
    }

    let output = Command::new("git")
        .args([
            "clone",
            "--depth",
            "1",
            "https://github.com/libjxl/conformance",
        ])
        .arg(dir)
        .output()
        .map_err(|e| format!("Failed to run git clone: {}", e))?;

    if !output.status.success() {
        return Err(format!(
            "git clone failed: {}",
            String::from_utf8_lossy(&output.stderr)
        ));
    }

    // Download reference data
    eprintln!("Downloading conformance reference data (this may take a few minutes)...");
    run_download_script(dir)
}

/// Run the download script to fetch reference NPY files.
fn run_download_script(dir: &Path) -> Result<(), String> {
    let script = dir
        .join("scripts")
        .join("download_and_symlink_using_curl.sh");

    if !script.exists() {
        return Err(format!("Download script not found: {:?}", script));
    }

    let output = Command::new("bash")
        .arg(&script)
        .current_dir(dir)
        .output()
        .map_err(|e| format!("Failed to run download script: {}", e))?;

    if !output.status.success() {
        return Err(format!(
            "Download script failed: {}",
            String::from_utf8_lossy(&output.stderr)
        ));
    }

    eprintln!("Conformance data ready.");
    Ok(())
}

/// Metadata for a conformance test frame.
#[derive(Debug, Clone)]
struct FrameInfo {
    #[allow(dead_code)]
    name: String,
    rms_error: f32,
    peak_error: f32,
}

/// Metadata for a conformance test case.
#[derive(Debug)]
struct ConformanceTestCase {
    name: String,
    path: PathBuf,
    frames: Vec<FrameInfo>,
    #[allow(dead_code)]
    intensity_target: f32,
    #[allow(dead_code)]
    extra_channel_types: Vec<String>,
    #[allow(dead_code)]
    bits_per_sample: Vec<u32>,
}

impl ConformanceTestCase {
    fn input_jxl(&self) -> PathBuf {
        self.path.join("input.jxl")
    }

    fn reference_npy(&self) -> PathBuf {
        self.path.join("reference_image.npy")
    }

    /// Get the reference ICC profile if it exists.
    /// This is the target color space for the reference output.
    fn reference_icc(&self) -> Option<PathBuf> {
        let path = self.path.join("reference.icc");
        if path.exists() { Some(path) } else { None }
    }

    /// Get the original ICC profile if it exists.
    /// This is the color space of the decoded image.
    fn original_icc(&self) -> Option<PathBuf> {
        let path = self.path.join("original.icc");
        if path.exists() { Some(path) } else { None }
    }
}

/// Parse test.json for a conformance test case.
fn parse_test_json(path: &Path) -> Option<ConformanceTestCase> {
    let test_json_path = path.join("test.json");
    let content = std::fs::read_to_string(&test_json_path).ok()?;
    let json: serde_json::Value = serde_json::from_str(&content).ok()?;

    let frames: Vec<FrameInfo> = json
        .get("frames")?
        .as_array()?
        .iter()
        .map(|f| FrameInfo {
            name: f
                .get("name")
                .and_then(|n| n.as_str())
                .unwrap_or("")
                .to_string(),
            rms_error: f.get("rms_error").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32,
            peak_error: f.get("peak_error").and_then(|v| v.as_f64()).unwrap_or(0.0) as f32,
        })
        .collect();

    let intensity_target = json
        .get("intensity_target")
        .and_then(|v| v.as_f64())
        .unwrap_or(255.0) as f32;

    let extra_channel_types: Vec<String> = json
        .get("extra_channel_type")
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(String::from))
                .collect()
        })
        .unwrap_or_default();

    let bits_per_sample: Vec<u32> = json
        .get("bits_per_sample")
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_u64().map(|n| n as u32))
                .collect()
        })
        .unwrap_or_default();

    Some(ConformanceTestCase {
        name: path.file_name()?.to_str()?.to_string(),
        path: path.to_path_buf(),
        frames,
        intensity_target,
        extra_channel_types,
        bits_per_sample,
    })
}

/// Discover all conformance test cases.
fn discover_conformance_tests() -> Vec<ConformanceTestCase> {
    // Ensure conformance data is available (auto-fetch if needed)
    let conformance_dir = match ensure_conformance_data() {
        Ok(dir) => dir,
        Err(e) => {
            eprintln!("Conformance data not available: {}", e);
            return vec![];
        }
    };

    let testcases_dir = conformance_dir.join("testcases");

    // Read main_level5.txt to get the list of tests
    let corpus_path = testcases_dir.join("main_level5.txt");
    let test_names: Vec<String> = if corpus_path.exists() {
        std::fs::read_to_string(&corpus_path)
            .unwrap_or_default()
            .lines()
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
            .map(|s| s.trim().to_string())
            .collect()
    } else {
        // Fall back to listing directories
        std::fs::read_dir(&testcases_dir)
            .ok()
            .map(|entries| {
                entries
                    .filter_map(|e| e.ok())
                    .filter(|e| e.path().is_dir())
                    .filter_map(|e| e.file_name().into_string().ok())
                    .collect()
            })
            .unwrap_or_default()
    };

    test_names
        .into_iter()
        .filter_map(|name| {
            let test_path = testcases_dir.join(&name);
            if test_path.exists() {
                parse_test_json(&test_path)
            } else {
                None
            }
        })
        .collect()
}

/// Simple NPY file reader for f32 arrays.
/// NPY format: https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html
fn read_npy_f32(path: &Path) -> Result<(Vec<usize>, Vec<f32>), String> {
    let data = std::fs::read(path).map_err(|e| format!("Failed to read NPY: {}", e))?;

    // Check magic number
    if data.len() < 10 || &data[0..6] != b"\x93NUMPY" {
        return Err("Invalid NPY magic number".to_string());
    }

    let major_version = data[6];
    let _minor_version = data[7];

    // Parse header length
    let header_len = if major_version == 1 {
        u16::from_le_bytes([data[8], data[9]]) as usize
    } else {
        u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize
    };

    let header_start = if major_version == 1 { 10 } else { 12 };
    let header_end = header_start + header_len;
    let header =
        std::str::from_utf8(&data[header_start..header_end]).map_err(|_| "Invalid NPY header")?;

    // Parse shape from header like "{'descr': '<f4', 'fortran_order': False, 'shape': (1, 1024, 1024, 4), }"
    let shape_start = header.find("'shape': (").ok_or("No shape in NPY header")?;
    let shape_str_start = shape_start + "'shape': (".len();
    let shape_str_end = header[shape_str_start..]
        .find(')')
        .ok_or("Invalid shape in NPY header")?
        + shape_str_start;
    let shape_str = &header[shape_str_start..shape_str_end];

    let shape: Vec<usize> = shape_str
        .split(',')
        .filter_map(|s| s.trim().parse().ok())
        .collect();

    // Verify dtype is f32 (little-endian)
    if !header.contains("'<f4'") && !header.contains("'float32'") {
        return Err(format!("Unsupported NPY dtype (expected f32): {}", header));
    }

    // Read data
    let data_start = header_end;
    let num_elements: usize = shape.iter().product();
    let expected_bytes = num_elements * 4;

    if data.len() < data_start + expected_bytes {
        return Err(format!(
            "NPY file too short: {} < {}",
            data.len(),
            data_start + expected_bytes
        ));
    }

    let pixels: Vec<f32> = data[data_start..data_start + expected_bytes]
        .chunks_exact(4)
        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .collect();

    Ok((shape, pixels))
}

/// Result of decoding a JXL file, including the output color profile.
#[cfg(feature = "cms")]
struct DecodeResult {
    frames: usize,
    height: usize,
    width: usize,
    channels: usize,
    pixels: Vec<f32>,
    output_profile: JxlColorProfile,
}

/// Decode a JXL file to f32 pixels for conformance testing.
/// Supports both single-frame and animated (multi-frame) images.
/// Pixels are returned as a flat buffer with shape [N, H, W, C].
///
/// All extra channels (alpha, spot colors, depth, etc.) are output as separate
/// planar buffers and interleaved into the final pixel array. This ensures the
/// output matches the NPY reference layout which includes all channels.
#[cfg(feature = "cms")]
fn decode_jxl_to_f32(path: &Path) -> Result<DecodeResult, String> {
    use crate::image::{Image, Rect};

    let data = std::fs::read(path).map_err(|e| format!("Failed to read JXL: {}", e))?;
    let mut input = data.as_slice();

    // Don't render spot colors into RGB — conformance references have raw channel data
    let options = JxlDecoderOptions {
        cms: Some(Box::new(MoxCms::new())),
        render_spot_colors: false,
        ..JxlDecoderOptions::default()
    };

    let mut decoder = JxlDecoder::<states::Initialized>::new(options);

    // Advance to image info
    let mut decoder_with_image_info = loop {
        match decoder.process(&mut input) {
            Ok(ProcessingResult::Complete { result }) => break result,
            Ok(ProcessingResult::NeedsMoreInput { fallback, .. }) => {
                if input.is_empty() {
                    return Err("Unexpected end of input during header".to_string());
                }
                decoder = fallback;
            }
            Err(e) => return Err(format!("Header decode error: {:?}", e)),
        }
    };

    let basic_info = decoder_with_image_info.basic_info().clone();
    let (width, height) = basic_info.size;

    // Use RGB/Grayscale as color type (no alpha in color buffer).
    // All extra channels (including alpha) get separate planar buffers
    // so we can interleave them to match the NPY reference layout.
    let default_format = decoder_with_image_info.current_pixel_format();
    let is_grayscale = matches!(
        default_format.color_type,
        JxlColorType::Grayscale | JxlColorType::GrayscaleAlpha
    );

    let color_type = if is_grayscale {
        JxlColorType::Grayscale
    } else {
        JxlColorType::Rgb
    };
    let color_channels = color_type.samples_per_pixel();

    // Request f32 output for color + all extra channels
    let num_extra_channels = basic_info.extra_channels.len();
    let total_channels = color_channels + num_extra_channels;

    let format = JxlPixelFormat {
        color_type,
        color_data_format: Some(JxlDataFormat::f32()),
        extra_channel_format: vec![Some(JxlDataFormat::f32()); num_extra_channels],
    };

    decoder_with_image_info.set_pixel_format(format);

    // Capture output color profile before advancing to frame info
    let output_profile = decoder_with_image_info.output_color_profile().clone();

    let frame_pixels = width * height * total_channels;
    let mut all_pixels = Vec::new();
    let mut frame_count = 0usize;

    // Multi-frame decode loop
    loop {
        // Advance to frame info
        let mut decoder_with_frame_info = loop {
            match decoder_with_image_info.process(&mut input) {
                Ok(ProcessingResult::Complete { result }) => break result,
                Ok(ProcessingResult::NeedsMoreInput { fallback, .. }) => {
                    if input.is_empty() {
                        return Err("Unexpected end of input before frame".to_string());
                    }
                    decoder_with_image_info = fallback;
                }
                Err(e) => return Err(format!("Frame info decode error: {:?}", e)),
            }
        };

        // Create output buffers: buffer 0 = interleaved color, buffers 1..N = planar extra channels
        let mut color_image = Image::<f32>::new((width * color_channels, height))
            .map_err(|e| format!("Color buffer error: {:?}", e))?;
        let mut extra_images: Vec<Image<f32>> = (0..num_extra_channels)
            .map(|_| Image::<f32>::new((width, height)))
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(|e| format!("Extra channel buffer error: {:?}", e))?;

        let mut buffers = Vec::with_capacity(1 + num_extra_channels);
        buffers.push(JxlOutputBuffer::from_image_rect_mut(
            color_image
                .get_rect_mut(Rect {
                    origin: (0, 0),
                    size: (width * color_channels, height),
                })
                .into_raw(),
        ));
        for ec_image in &mut extra_images {
            buffers.push(JxlOutputBuffer::from_image_rect_mut(
                ec_image
                    .get_rect_mut(Rect {
                        origin: (0, 0),
                        size: (width, height),
                    })
                    .into_raw(),
            ));
        }

        // Decode frame pixels
        decoder_with_image_info = loop {
            match decoder_with_frame_info.process(&mut input, &mut buffers) {
                Ok(ProcessingResult::Complete { result }) => break result,
                Ok(ProcessingResult::NeedsMoreInput { fallback, .. }) => {
                    if input.is_empty() {
                        return Err("Unexpected end of input during frame".to_string());
                    }
                    decoder_with_frame_info = fallback;
                }
                Err(e) => return Err(format!("Frame decode error: {:?}", e)),
            }
        };

        // Interleave color + extra channels into [C0, C1, C2, EC0, EC1, ...] per pixel
        all_pixels.reserve(frame_pixels);
        for y in 0..height {
            let color_row = color_image.row(y);
            for x in 0..width {
                // Color channels (interleaved in color_row)
                for c in 0..color_channels {
                    all_pixels.push(color_row[x * color_channels + c]);
                }
                // Extra channels (each is planar)
                for ec_image in &extra_images {
                    all_pixels.push(ec_image.row(y)[x]);
                }
            }
        }
        frame_count += 1;

        // Check for more frames
        if !decoder_with_image_info.has_more_frames() {
            break;
        }
    }

    Ok(DecodeResult {
        frames: frame_count,
        height,
        width,
        channels: total_channels,
        pixels: all_pixels,
        output_profile,
    })
}

/// Transform pixels from source color profile to target ICC profile.
/// Returns (output_channels, transformed_pixels).
#[cfg(feature = "cms")]
fn transform_pixels_to_icc(
    pixels: &[f32],
    width: usize,
    height: usize,
    channels: usize,
    source_profile: &JxlColorProfile,
    target_icc: &[u8],
    preserve_alpha: bool, // If true, preserve alpha channel when transforming RGBA to RGB
) -> Result<(usize, Vec<f32>), String> {
    let cms = MoxCms::new();

    let target_profile = JxlColorProfile::Icc(target_icc.to_vec());

    // Initialize transform
    let (cms_output_channels, mut transforms) = cms
        .initialize_transforms(
            1,
            width * height,
            source_profile.clone(),
            target_profile,
            255.0,
        )
        .map_err(|e| format!("Failed to create transform: {:?}", e))?;

    eprintln!(
        "    Transform: {} channels -> {} channels",
        channels, cms_output_channels
    );

    if transforms.is_empty() {
        return Err("No transform created".to_string());
    }

    let transform = &mut transforms[0];

    // Handle channel count differences
    // Source may be RGBA, target may be RGB (or vice versa)
    let input_channels = channels;
    let has_alpha = input_channels == 4 || input_channels == 2;

    // Determine actual output channels (may differ from CMS output if we preserve alpha)
    let final_output_channels = if has_alpha && cms_output_channels == 3 && preserve_alpha {
        4 // Keep alpha channel
    } else if has_alpha && cms_output_channels == 1 && preserve_alpha && input_channels == 2 {
        2 // Keep alpha for grayscale
    } else {
        cms_output_channels
    };

    // Allocate output buffer
    let mut output = vec![0.0f32; width * height * final_output_channels];

    // Transform row by row
    let row_pixels = width * input_channels;

    for y in 0..height {
        let input_start = y * row_pixels;

        // If we have alpha and CMS only outputs RGB, transform RGB and preserve alpha
        if input_channels == 4 && cms_output_channels == 3 && preserve_alpha {
            // Transform RGB, preserve alpha
            let mut input_rgb: Vec<f32> = Vec::with_capacity(width * 3);
            for x in 0..width {
                let idx = input_start + x * 4;
                input_rgb.push(pixels[idx]);
                input_rgb.push(pixels[idx + 1]);
                input_rgb.push(pixels[idx + 2]);
            }
            let mut output_rgb = vec![0.0f32; width * 3];
            transform
                .do_transform(&input_rgb, &mut output_rgb)
                .map_err(|e| format!("Transform failed: {:?}", e))?;

            // Write RGBA output
            let output_start = y * width * 4;
            for x in 0..width {
                output[output_start + x * 4] = output_rgb[x * 3];
                output[output_start + x * 4 + 1] = output_rgb[x * 3 + 1];
                output[output_start + x * 4 + 2] = output_rgb[x * 3 + 2];
                output[output_start + x * 4 + 3] = pixels[input_start + x * 4 + 3]; // preserve alpha
            }
        } else if input_channels == 4 && cms_output_channels == 3 && !preserve_alpha {
            // Strip alpha, output RGB
            let mut input_rgb: Vec<f32> = Vec::with_capacity(width * 3);
            for x in 0..width {
                let idx = input_start + x * 4;
                input_rgb.push(pixels[idx]);
                input_rgb.push(pixels[idx + 1]);
                input_rgb.push(pixels[idx + 2]);
            }
            let output_start = y * width * 3;
            transform
                .do_transform(
                    &input_rgb,
                    &mut output[output_start..output_start + width * 3],
                )
                .map_err(|e| format!("Transform failed: {:?}", e))?;
        } else if input_channels == 3 && cms_output_channels == 3 {
            // RGB to RGB - straightforward
            let output_start = y * width * 3;
            transform
                .do_transform(
                    &pixels[input_start..input_start + width * 3],
                    &mut output[output_start..output_start + width * 3],
                )
                .map_err(|e| format!("Transform failed: {:?}", e))?;
        } else if input_channels == 4 && cms_output_channels == 4 {
            // RGBA to RGBA
            let output_start = y * width * 4;
            transform
                .do_transform(
                    &pixels[input_start..input_start + width * 4],
                    &mut output[output_start..output_start + width * 4],
                )
                .map_err(|e| format!("Transform failed: {:?}", e))?;
        } else if input_channels == 1 && cms_output_channels == 1 {
            // Gray to Gray
            let output_start = y * width;
            transform
                .do_transform(
                    &pixels[input_start..input_start + width],
                    &mut output[output_start..output_start + width],
                )
                .map_err(|e| format!("Transform failed: {:?}", e))?;
        } else if input_channels == 2 && cms_output_channels == 1 && !preserve_alpha {
            // GrayAlpha to Gray - strip alpha
            let mut input_gray: Vec<f32> = Vec::with_capacity(width);
            for x in 0..width {
                input_gray.push(pixels[input_start + x * 2]);
            }
            let output_start = y * width;
            transform
                .do_transform(&input_gray, &mut output[output_start..output_start + width])
                .map_err(|e| format!("Transform failed: {:?}", e))?;
        } else if input_channels == 2 && cms_output_channels == 1 && preserve_alpha {
            // GrayAlpha to GrayAlpha - preserve alpha
            let mut input_gray: Vec<f32> = Vec::with_capacity(width);
            for x in 0..width {
                input_gray.push(pixels[input_start + x * 2]);
            }
            let mut output_gray = vec![0.0f32; width];
            transform
                .do_transform(&input_gray, &mut output_gray)
                .map_err(|e| format!("Transform failed: {:?}", e))?;

            let output_start = y * width * 2;
            for x in 0..width {
                output[output_start + x * 2] = output_gray[x];
                output[output_start + x * 2 + 1] = pixels[input_start + x * 2 + 1]; // preserve alpha
            }
        } else {
            return Err(format!(
                "Unsupported channel conversion: {} -> {}",
                input_channels, cms_output_channels
            ));
        }
    }

    Ok((final_output_channels, output))
}

/// Compute peak error between decoded and reference pixels.
/// This is a simple helper for debugging.
#[cfg(feature = "cms")]
fn compute_peak_error(
    decoded: &[f32],
    reference: &[f32],
    dec_channels: usize,
    ref_channels: usize,
) -> f32 {
    let compare_channels = dec_channels.min(ref_channels);
    let mut max_error: f32 = 0.0;

    let dec_stride = dec_channels;
    let ref_stride = ref_channels;

    let num_pixels = decoded.len() / dec_channels;
    let ref_pixels = reference.len() / ref_channels;
    let compare_pixels = num_pixels.min(ref_pixels);

    for i in 0..compare_pixels {
        let dec_base = i * dec_stride;
        let ref_base = i * ref_stride;

        for c in 0..compare_channels {
            let dec_val = decoded[dec_base + c];
            let ref_val = reference[ref_base + c];
            let error = (dec_val - ref_val).abs();
            max_error = max_error.max(error);
        }
    }

    max_error
}

/// Compare decoded output against NPY reference using conformance thresholds.
fn compare_conformance(
    test: &ConformanceTestCase,
    decoded_shape: (usize, usize, usize, usize), // (frames, height, width, channels)
    decoded: &[f32],
    ref_shape: &[usize],
    reference: &[f32],
) -> Result<(), String> {
    let (dec_frames, dec_height, dec_width, dec_channels) = decoded_shape;

    // Reference shape is (frames, height, width, channels)
    if ref_shape.len() != 4 {
        return Err(format!("Invalid reference shape: {:?}", ref_shape));
    }
    let (ref_frames, ref_height, ref_width, ref_channels) =
        (ref_shape[0], ref_shape[1], ref_shape[2], ref_shape[3]);

    // Allow decoded to have more frames (e.g., last frame repeated)
    if dec_frames < ref_frames {
        return Err(format!(
            "Frame count mismatch: decoded {} < reference {}",
            dec_frames, ref_frames
        ));
    }

    if dec_height != ref_height || dec_width != ref_width {
        return Err(format!(
            "Size mismatch: decoded {}x{} vs reference {}x{}",
            dec_width, dec_height, ref_width, ref_height
        ));
    }

    // Allow decoded RGBA when reference is RGB (if alpha is all 1.0)
    let compare_channels = if dec_channels == 4 && ref_channels == 3 {
        3 // Compare only RGB
    } else if dec_channels != ref_channels {
        return Err(format!(
            "Channel count mismatch: decoded {} vs reference {}",
            dec_channels, ref_channels
        ));
    } else {
        dec_channels
    };

    let frame_size_dec = dec_height * dec_width * dec_channels;
    let frame_size_ref = ref_height * ref_width * ref_channels;

    for frame_idx in 0..ref_frames {
        let frame_info = test
            .frames
            .get(frame_idx)
            .ok_or_else(|| format!("No frame info for frame {}", frame_idx))?;

        let dec_frame_start = frame_idx * frame_size_dec;
        let ref_frame_start = frame_idx * frame_size_ref;

        let mut max_error: f32 = 0.0;
        let mut sum_sq_errors: Vec<f64> = vec![0.0; compare_channels];
        let pixel_count = dec_height * dec_width;

        for y in 0..dec_height {
            for x in 0..dec_width {
                let dec_idx = dec_frame_start + (y * dec_width + x) * dec_channels;
                let ref_idx = ref_frame_start + (y * ref_width + x) * ref_channels;

                for c in 0..compare_channels {
                    let dec_val = decoded[dec_idx + c];
                    let ref_val = reference[ref_idx + c];
                    let error = (dec_val - ref_val).abs();

                    if error > max_error {
                        max_error = error;
                        if error > frame_info.peak_error {
                            eprintln!(
                                "    New max error {:.6} at ({}, {}, ch={}) dec={:.6} ref={:.6}",
                                error, x, y, c, dec_val, ref_val
                            );
                        }
                    }
                    sum_sq_errors[c] += (error as f64) * (error as f64);
                }
            }
        }
        // Dump pixels with large errors (all channels)
        if max_error > frame_info.peak_error {
            let mut error_pixels: Vec<(f32, usize, usize)> = Vec::new();
            for y in 0..dec_height {
                for x in 0..dec_width {
                    let dec_idx = dec_frame_start + (y * dec_width + x) * dec_channels;
                    let ref_idx = ref_frame_start + (y * ref_width + x) * ref_channels;
                    let mut pixel_error: f32 = 0.0;
                    for c in 0..compare_channels {
                        let e = (decoded[dec_idx + c] - reference[ref_idx + c]).abs();
                        pixel_error = pixel_error.max(e);
                    }
                    if pixel_error > frame_info.peak_error {
                        error_pixels.push((pixel_error, x, y));
                    }
                }
            }
            error_pixels.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
            eprintln!(
                "    {} pixels exceed threshold. Top 10:",
                error_pixels.len()
            );
            for &(err, x, y) in error_pixels.iter().take(10) {
                let di = dec_frame_start + (y * dec_width + x) * dec_channels;
                let ri = ref_frame_start + (y * ref_width + x) * ref_channels;
                eprintln!(
                    "    ({:3},{:3}) err={:.6} dec=[{:.8},{:.8},{:.8}] ref=[{:.8},{:.8},{:.8}]",
                    x,
                    y,
                    err,
                    decoded[di],
                    decoded[di + 1],
                    decoded[di + 2],
                    reference[ri],
                    reference[ri + 1],
                    reference[ri + 2]
                );
            }
        }

        // Compute per-channel RMSE and take max
        let max_rmse: f32 = sum_sq_errors
            .iter()
            .map(|&sum| (sum / pixel_count as f64).sqrt() as f32)
            .fold(0.0f32, f32::max);

        // Check against thresholds
        if max_error > frame_info.peak_error {
            return Err(format!(
                "Frame {}: peak error {} > threshold {}",
                frame_idx, max_error, frame_info.peak_error
            ));
        }

        if max_rmse > frame_info.rms_error {
            return Err(format!(
                "Frame {}: RMSE {} > threshold {}",
                frame_idx, max_rmse, frame_info.rms_error
            ));
        }

        eprintln!(
            "  Frame {}: peak_error={:.6} (limit {:.6}), rmse={:.6} (limit {:.6})",
            frame_idx, max_error, frame_info.peak_error, max_rmse, frame_info.rms_error
        );
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    /// List all available conformance tests.
    #[test]
    #[ignore] // Run with: CONFORMANCE_PATH=/path/to/conformance cargo test list_conformance_tests -- --ignored --nocapture
    fn list_conformance_tests() {
        let tests = discover_conformance_tests();
        if tests.is_empty() {
            eprintln!("No conformance tests found. Set CONFORMANCE_PATH environment variable.");
            eprintln!("Example: CONFORMANCE_PATH=~/work/conformance cargo test ...");
            return;
        }

        eprintln!("Found {} conformance tests:", tests.len());
        for test in &tests {
            eprintln!("  {} ({} frames)", test.name, test.frames.len());
        }
    }

    /// Run a single conformance test by name.
    #[cfg(feature = "cms")]
    fn run_conformance_test(name: &str) -> Result<(), String> {
        let tests = discover_conformance_tests();
        let test = tests
            .iter()
            .find(|t| t.name == name)
            .ok_or_else(|| format!("Test '{}' not found", name))?;

        let input_path = test.input_jxl();
        let ref_path = test.reference_npy();

        if !input_path.exists() {
            return Err(format!("Input file not found: {:?}", input_path));
        }
        if !ref_path.exists() {
            return Err(format!(
                "Reference NPY not found: {:?}. Run download_and_symlink_using_curl.sh",
                ref_path
            ));
        }

        eprintln!("Testing: {}", name);

        // Load reference
        let (ref_shape, reference) = read_npy_f32(&ref_path)?;
        eprintln!("  Reference shape: {:?}", ref_shape);

        // Decode JXL
        let result = decode_jxl_to_f32(&input_path)?;
        eprintln!(
            "  Decoded shape: ({}, {}, {}, {})",
            result.frames, result.height, result.width, result.channels
        );

        eprintln!("  Output color profile: {}", result.output_profile);

        // Debug: show detailed profile info
        if let JxlColorProfile::Simple(enc) = &result.output_profile {
            eprintln!("    Detailed: {}", enc.get_color_encoding_description());
        }
        if let Some(icc) = result.output_profile.try_as_icc() {
            eprintln!("    Generated ICC size: {} bytes", icc.len());
            // Compare with original.icc if available
            if let Some(orig_icc_path) = test.original_icc()
                && let Ok(orig_icc) = std::fs::read(&orig_icc_path)
            {
                if icc.as_ref() == &orig_icc {
                    eprintln!("    Generated ICC matches original.icc exactly");
                } else {
                    eprintln!(
                        "    Generated ICC differs from original.icc ({} bytes)",
                        orig_icc.len()
                    );
                }
            }
        }

        // Check if we need to transform to reference ICC
        let (final_pixels, final_channels) = if let Some(ref_icc_path) = test.reference_icc() {
            let ref_icc = std::fs::read(&ref_icc_path)
                .map_err(|e| format!("Failed to read reference.icc: {}", e))?;

            // Use output_color_profile for transformation:
            // The decoder outputs pixels in output_color_profile color space.
            // We need to transform from output_color_profile to reference.icc.
            // (Similar to how libjxl outputs decoded.icc via --icc_out)
            let source_profile = result.output_profile.clone();
            eprintln!("  Source profile for transform: {}", source_profile);

            // Check if source profile matches reference ICC
            let source_icc = source_profile.try_as_icc();
            let needs_transform = match &source_icc {
                Some(src_icc) => {
                    let differs = src_icc.as_ref() != &ref_icc;
                    if differs {
                        eprintln!(
                            "  Source ICC ({} bytes) differs from reference ICC ({} bytes)",
                            src_icc.len(),
                            ref_icc.len()
                        );
                    }
                    differs
                }
                None => true, // Can't compare, assume different
            };

            if needs_transform {
                eprintln!("  Reference ICC size: {} bytes", ref_icc.len());

                // Compute error without transform to check if raw decode is already good
                let ref_channels = ref_shape.get(3).copied().unwrap_or(result.channels);
                let raw_error =
                    compute_peak_error(&result.pixels, &reference, result.channels, ref_channels);
                eprintln!("  Raw decode peak error (no transform): {:.6}", raw_error);

                // Get the threshold for this frame
                let threshold = test.frames.first().map(|f| f.peak_error).unwrap_or(0.06);

                // If raw decode is within threshold, skip transform
                // This handles the case where our ICC differs from libjxl's ICC
                // but both describe equivalent color spaces
                if raw_error <= threshold {
                    eprintln!("  Raw decode within threshold, skipping transform");
                    (result.pixels, result.channels)
                } else {
                    eprintln!("  Transforming to reference ICC...");

                    // Debug: show sample input values before transform
                    let mid_pixel = result.width * result.height / 2;
                    let mid_idx = mid_pixel * result.channels;
                    eprintln!(
                        "  Sample input pixel (mid): {:?}",
                        &result.pixels[mid_idx..mid_idx + result.channels.min(4)]
                    );

                    // Preserve alpha if reference has alpha
                    let preserve_alpha = ref_channels == 4 || ref_channels == 2;
                    // For multi-frame, treat as a taller image (height * frames)
                    // since transform_pixels_to_icc processes row-by-row
                    let total_height = result.height * result.frames;
                    let (trans_channels, transformed) = transform_pixels_to_icc(
                        &result.pixels,
                        result.width,
                        total_height,
                        result.channels,
                        &source_profile,
                        &ref_icc,
                        preserve_alpha,
                    )?;

                    // Debug: show sample output values after transform
                    let mid_out_idx = mid_pixel * trans_channels;
                    eprintln!(
                        "  Sample output pixel (mid): {:?}",
                        &transformed[mid_out_idx..mid_out_idx + trans_channels.min(4)]
                    );
                    // Also show reference at same location
                    let mid_ref_idx = mid_pixel * ref_channels;
                    eprintln!(
                        "  Sample reference pixel (mid): {:?}",
                        &reference[mid_ref_idx..mid_ref_idx + ref_channels.min(4)]
                    );

                    // Compute error after transform
                    let transformed_error =
                        compute_peak_error(&transformed, &reference, trans_channels, ref_channels);
                    eprintln!("  Transformed peak error: {:.6}", transformed_error);

                    // Use whichever gives better results
                    if raw_error <= transformed_error {
                        eprintln!("  Raw decode better than transformed, using raw");
                        (result.pixels, result.channels)
                    } else {
                        eprintln!("  Using transformed result");
                        (transformed, trans_channels)
                    }
                }
            } else {
                eprintln!("  Output profile matches reference ICC, no transform needed");
                (result.pixels, result.channels)
            }
        } else {
            // No reference ICC, use decoded pixels directly
            (result.pixels, result.channels)
        };

        // Compare
        compare_conformance(
            test,
            (result.frames, result.height, result.width, final_channels),
            &final_pixels,
            &ref_shape,
            &reference,
        )?;

        eprintln!("  PASS");
        Ok(())
    }

    #[cfg(not(feature = "cms"))]
    fn run_conformance_test(_name: &str) -> Result<(), String> {
        Err("CMS feature required for conformance tests".to_string())
    }

    /// Run all conformance tests.
    #[test]
    #[ignore] // Run with: CONFORMANCE_PATH=/path/to/conformance cargo test --features cms run_all_conformance -- --ignored --nocapture
    fn run_all_conformance() {
        let tests = discover_conformance_tests();
        if tests.is_empty() {
            eprintln!("No conformance tests found. Set CONFORMANCE_PATH.");
            eprintln!("1. git clone https://github.com/libjxl/conformance");
            eprintln!("2. cd conformance && bash scripts/download_and_symlink_using_curl.sh");
            eprintln!(
                "3. CONFORMANCE_PATH=/path/to/conformance cargo test --features cms run_all_conformance -- --ignored --nocapture"
            );
            return;
        }

        eprintln!("Running {} conformance tests...\n", tests.len());

        let mut passed = 0;
        let mut failed = 0;
        let mut failures: Vec<(String, String)> = Vec::new();

        for test in &tests {
            match run_conformance_test(&test.name) {
                Ok(()) => passed += 1,
                Err(e) => {
                    failed += 1;
                    failures.push((test.name.clone(), e));
                }
            }
            eprintln!();
        }

        eprintln!("\n=== Results ===");
        eprintln!("Passed: {}/{}", passed, passed + failed);
        eprintln!("Failed: {}/{}", failed, passed + failed);

        if !failures.is_empty() {
            eprintln!("\nFailures:");
            for (name, error) in &failures {
                eprintln!("  {}: {}", name, error);
            }
            panic!("{} conformance tests failed", failed);
        }
    }

    // Individual test functions for each conformance test.
    // These allow running specific tests and seeing which ones pass in CI.

    macro_rules! conformance_test {
        ($name:ident) => {
            #[test]
            #[ignore]
            fn $name() {
                if conformance_dir().is_none() {
                    eprintln!("Skipped: CONFORMANCE_PATH not set");
                    return;
                }
                run_conformance_test(stringify!($name)).unwrap();
            }
        };
    }

    // Level 5 conformance tests
    conformance_test!(alpha_nonpremultiplied);
    conformance_test!(alpha_triangles);
    conformance_test!(animation_icos4d_5);
    conformance_test!(animation_newtons_cradle);
    conformance_test!(animation_spline_5);
    conformance_test!(bench_oriented_brg_5);
    conformance_test!(bicycles);
    conformance_test!(bike_5);
    conformance_test!(blendmodes_5);
    conformance_test!(cafe_5);
    conformance_test!(delta_palette);
    conformance_test!(grayscale_5);
    conformance_test!(grayscale_jpeg_5);
    conformance_test!(grayscale_public_university);
    conformance_test!(lz77_flower);
    conformance_test!(noise_5);
    conformance_test!(opsin_inverse_5);
    conformance_test!(patches_5);
    conformance_test!(patches_lossless);
    conformance_test!(progressive_5);
    conformance_test!(spot);
    conformance_test!(sunset_logo);
    conformance_test!(upsampling_5);
}