use std::{io::Cursor, mem::MaybeUninit};

use image::{codecs::png::PngDecoder, ImageDecoder};

#[repr(C)]
pub struct ImageInfo {
    width: u32,
    height: u32,
    byte_count: u64,
    color_type: ColorType,
}

#[repr(u8)]
pub enum ColorType {
    L8,
    La8,
    Rgb8,
    Rgba8,
    L16,
    La16,
    Rgb16,
    Rgba16,
    Rgb32F,
    Rgba32F,
}

impl ColorType {
    fn from(value: image::ColorType) -> Option<Self> {
        Some(match value {
            image::ColorType::L8 => Self::L8,
            image::ColorType::La8 => Self::La8,
            image::ColorType::Rgb8 => Self::Rgb8,
            image::ColorType::Rgba8 => Self::Rgba8,
            image::ColorType::L16 => Self::L16,
            image::ColorType::La16 => Self::La16,
            image::ColorType::Rgb16 => Self::Rgb16,
            image::ColorType::Rgba16 => Self::Rgba16,
            image::ColorType::Rgb32F => Self::Rgb32F,
            image::ColorType::Rgba32F => Self::Rgba32F,
            _ => return None,
        })
    }
}

/// Returns `true` to indicate failure.
#[no_mangle]
pub extern "system" fn load_png_info(input: *const u8, input_len: usize, result: &mut MaybeUninit<ImageInfo>) -> bool {
    let input = unsafe { std::slice::from_raw_parts(input, input_len) };
    let Ok(dec) = PngDecoder::new(Cursor::new(input)) else {
        return true;
    };
    let (width, height) = dec.dimensions();
    let Some(color_type) = ColorType::from(dec.color_type()) else {
        return true;
    };
    result.write(ImageInfo {
        width,
        height,
        byte_count: dec.total_bytes(),
        color_type,
    });
    false
}

/// Returns `true` to indicate failure.
#[no_mangle]
pub extern "system" fn load_png(input: *const u8, input_len: usize, output: *mut u8, output_len: usize) -> bool {
    let input = unsafe { std::slice::from_raw_parts(input, input_len) };
    let output = unsafe { std::slice::from_raw_parts_mut(output, output_len) };
    let Ok(dec) = PngDecoder::new(Cursor::new(input)) else {
        return true;
    };
    let Ok(byte_count) = usize::try_from(dec.total_bytes()) else {
        return true;
    };
    if byte_count != output.len() {
        return true;
    }
    if let Err(_) = dec.read_image(output) {
        return true;
    }
    false
}