//! Generate reference values for cpp_parity_locked tests.
//!
//! Run with: cargo run --release --example generate_parity_values

use butteraugli::ButteraugliParams;
use dssim::Dssim;
use jpegli::decode::Decoder;
use jpegli::types::Subsampling;
use jpegli::{Encoder, PixelFormat, Quality};
use rgb::RGBA8;
use std::fs;

fn rgb_to_rgba(data: &[u8]) -> Vec<RGBA8> {
    data.chunks(3)
        .map(|c| RGBA8::new(c[0], c[1], c[2], 255))
        .collect()
}

fn main() {
    let png_path = std::path::Path::new("tests/images/1.png");
    let png_data = fs::read(png_path).expect("Failed to read test image");
    let decoder = png::Decoder::new(&png_data[..]);
    let mut reader = decoder.read_info().unwrap();
    let mut buf = vec![0; reader.output_buffer_size()];
    let info = reader.next_frame(&mut buf).unwrap();

    let rgb: Vec<u8> = match info.color_type {
        png::ColorType::Rgb => buf[..info.buffer_size()].to_vec(),
        png::ColorType::Rgba => buf[..info.buffer_size()]
            .chunks(4)
            .flat_map(|c| [c[0], c[1], c[2]])
            .collect(),
        _ => panic!("Unsupported color type"),
    };

    let width = info.width as usize;
    let height = info.height as usize;

    println!("// Test image: tests/images/1.png ({}x{})", width, height);
    println!("// Generated by: cargo run --release --example generate_parity_values");
    println!("// Format: (quality, size, dssim, butteraugli)");
    println!();

    // Quality levels
    let qualities: Vec<u8> = (1..=20).map(|i| i * 5).collect();

    // Configurations to test
    let configs = [
        (Subsampling::S444, true, "S444_OPT"),
        (Subsampling::S444, false, "S444_FIXED"),
        (Subsampling::S420, true, "S420_OPT"),
        (Subsampling::S422, true, "S422_OPT"),
    ];

    // Prepare DSSIM
    let dssim = Dssim::new();
    let rgba_orig = rgb_to_rgba(&rgb);
    let orig_img = dssim.create_image_rgba(&rgba_orig, width, height).unwrap();

    // Butteraugli params
    let bfly_params = ButteraugliParams::default();

    for (subsampling, optimize, name) in configs {
        println!("/// {} - (quality, size, dssim, butteraugli)", name);
        println!("const RUST_{}: &[(u8, usize, f64, f64)] = &[", name);

        for &q in &qualities {
            let jpeg = Encoder::new()
                .width(width as u32)
                .height(height as u32)
                .pixel_format(PixelFormat::Rgb)
                .quality(Quality::from_quality(q as f32))
                .subsampling(subsampling)
                .optimize_huffman(optimize)
                .encode(&rgb)
                .expect("Encoding failed");

            let size = jpeg.len();

            // Decode
            let decoded = Decoder::new()
                .output_format(PixelFormat::Rgb)
                .decode(&jpeg)
                .expect("Decoding failed");

            // DSSIM
            let rgba_dec = rgb_to_rgba(&decoded.data);
            let dec_img = dssim.create_image_rgba(&rgba_dec, width, height).unwrap();
            let (dssim_val, _) = dssim.compare(&orig_img, dec_img);

            // Butteraugli
            let bfly =
                butteraugli::compute_butteraugli(&rgb, &decoded.data, width, height, &bfly_params)
                    .expect("butteraugli")
                    .score;

            println!(
                "    ({}, {}, {:.6}, {:.4}),",
                q,
                size,
                f64::from(dssim_val),
                bfly
            );
        }

        println!("];");
        println!();
    }
}
