palette-bin 0.1.0

use image::{ImageBuffer, Rgb, RgbImage};
use palette::color_difference::EuclideanDistance;
use palette::{IntoColor, Lab, Srgb};
use std::cmp::Ordering;
use std::fs::File;
use std::io;
use std::io::Read;
use std::path::PathBuf;

/// Parse an ACT (Adobe Color Table) file and return the contained colors
fn parse_act_palette(act_path: &PathBuf) -> Result<Vec<Rgb<u8>>, io::Error> {
    let mut file = File::open(act_path)?;
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer)?;

    let mut num_colors = 256;

    // Check if the file has the color count at the end (file size >= 768 + 2)
    if buffer.len() >= 770 {
        let count_bytes = [buffer[768], buffer[769]];
        let specified_count = u16::from_be_bytes(count_bytes) as usize;

        // Only use the specified count if it's valid (> 0 and <= 256)
        if specified_count > 0 && specified_count <= 256 {
            num_colors = specified_count;
        }
    }

    // Parse the colors
    let mut colors = Vec::with_capacity(num_colors);
    for i in 0..num_colors {
        if i * 3 + 2 < buffer.len() {
            let r = buffer[i * 3];
            let g = buffer[i * 3 + 1];
            let b = buffer[i * 3 + 2];
            colors.push(Rgb([r, g, b]));
        } else {
            break; // Stop if we run out of data
        }
    }

    // Verify that we got the expected number of colors
    if colors.len() != num_colors {
        println!(
            "Warning: Expected {} colors but got {}",
            num_colors,
            colors.len()
        );
    }

    Ok(colors)
}

/// Applies an ACT palette to an image using Floyd-Steinberg dithering.
pub fn apply_palette(
    img: RgbImage,
    act_path: &PathBuf,
) -> Result<RgbImage, Box<dyn std::error::Error>> {
    let palette = parse_act_palette(act_path)?;
    let (width, height) = img.dimensions();

    let mut buffer = img.clone();
    let mut result = ImageBuffer::from_pixel(width, height, Rgb([255, 255, 255]));

    for y in 0..height {
        for x in 0..width {
            let old_pixel = *buffer.get_pixel(x, y);

            let new_pixel = find_closest_color(&old_pixel, &palette);
            result.put_pixel(x, y, new_pixel);

            let error_r = old_pixel[0] as i32 - new_pixel[0] as i32;
            let error_g = old_pixel[1] as i32 - new_pixel[1] as i32;
            let error_b = old_pixel[2] as i32 - new_pixel[2] as i32;

            // Distribute the error to neighboring pixels using Floyd-Steinberg algorithm
            // Weights: 7/16 to the right, 3/16 to the bottom-left, 5/16 to the bottom, 1/16 to the bottom-right

            // Right pixel (x+1, y)
            if x + 1 < width {
                let pixel = buffer.get_pixel_mut(x + 1, y);
                distribute_error(pixel, error_r, error_g, error_b, 7.0 / 16.0);
            }

            // Next row pixels
            if y + 1 < height {
                // Bottom-left pixel (x-1, y+1)
                if x > 0 {
                    let pixel = buffer.get_pixel_mut(x - 1, y + 1);
                    distribute_error(pixel, error_r, error_g, error_b, 3.0 / 16.0);
                }

                // Bottom pixel (x, y+1)
                let pixel = buffer.get_pixel_mut(x, y + 1);
                distribute_error(pixel, error_r, error_g, error_b, 5.0 / 16.0);

                // Bottom-right pixel (x+1, y+1)
                if x + 1 < width {
                    let pixel = buffer.get_pixel_mut(x + 1, y + 1);
                    distribute_error(pixel, error_r, error_g, error_b, 1.0 / 16.0);
                }
            }
        }
    }

    Ok(result)
}

/// Helper function to distribute error to neighboring pixels
fn distribute_error(pixel: &mut Rgb<u8>, error_r: i32, error_g: i32, error_b: i32, weight: f32) {
    // Calculate new values with error diffusion
    let r = (pixel[0] as f32 + error_r as f32 * weight).clamp(0.0, 255.0) as u8;
    let g = (pixel[1] as f32 + error_g as f32 * weight).clamp(0.0, 255.0) as u8;
    let b = (pixel[2] as f32 + error_b as f32 * weight).clamp(0.0, 255.0) as u8;

    *pixel = Rgb([r, g, b]);
}

/// Find the index of the closest color in the palette to the given color
fn find_closest_color(source: &Rgb<u8>, palette: &[Rgb<u8>]) -> Rgb<u8> {
    let source_rgb = Srgb::new(
        source[0] as f32 / 255.0,
        source[1] as f32 / 255.0,
        source[2] as f32 / 255.0,
    );
    let source_lab: Lab = source_rgb.into_color();

    palette
        .iter()
        .min_by(|&a, &b| {
            let a_rgb = Srgb::new(
                a[0] as f32 / 255.0,
                a[1] as f32 / 255.0,
                a[2] as f32 / 255.0,
            );
            let a_lab: Lab = a_rgb.into_color();

            let b_rgb = Srgb::new(
                b[0] as f32 / 255.0,
                b[1] as f32 / 255.0,
                b[2] as f32 / 255.0,
            );
            let b_lab: Lab = b_rgb.into_color();

            let dist_a = source_lab.distance_squared(a_lab);
            let dist_b = source_lab.distance_squared(b_lab);

            dist_a.partial_cmp(&dist_b).unwrap_or(Ordering::Equal)
        })
        .cloned()
        .unwrap_or(*source)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{create_test_act_file, create_test_image};
    use tempfile::tempdir;

    #[test]
    fn test_apply_palette() {
        let img = create_test_image(2, 2);
        let temp_dir = tempdir().unwrap();
        let act_path = temp_dir.path().join("test_palette.act");
        create_test_act_file(&act_path).unwrap();

        let quantized_img = apply_palette(img, &act_path).unwrap();

        assert_eq!(quantized_img.get_pixel(0, 0), &Rgb([0, 0, 0]));
        assert_eq!(quantized_img.get_pixel(0, 1), &Rgb([0, 255, 0]));
        assert_eq!(quantized_img.get_pixel(1, 0), &Rgb([255, 0, 0]));
        assert_eq!(quantized_img.get_pixel(1, 1), &Rgb([255, 255, 255]));
        std::fs::remove_file(act_path).unwrap();
    }
}