//! Image codec wrappers. PNG (universal floor), JPEG (compatibility),
//! AVIF (modern lossy, AV1-based), JPEG-XL (lossless + lossy). A build with the
//! `floor` feature alone has PNG, through the `image` crate's pure-Rust codec; the
//! others need `codecs`.
//!
//! Each codec exposes `encode(rgb, h, w, quality) -> Vec<u8>` and
//! `decode(bytes) -> (rgb, h, w)`. RGB is row-major, 8-bit, 3 channels.
//! `quality` is in 1..=100 where applicable.
//!
//! Every decoder holds a picture to the reader limits of SPEC §2: at most
//! [`IMAGE_MAX_SIDE`] wide and tall and at most [`IMAGE_MAX_PIXELS`] pixels. It refuses
//! one past them (`image_too_large`) from the dimensions the payload declares, before it
//! decodes a pixel, and decodes only the format its capability names.

use std::io::Cursor;

use image::ImageEncoder;
#[cfg(feature = "codecs")]
use jpegxl_rs::encode::EncoderSpeed;
#[cfg(feature = "codecs")]
use jpegxl_rs::{decoder_builder, encoder_builder};

#[derive(Debug)]
pub struct ImageError(pub String);

impl<E: std::fmt::Display> From<E> for ImageError {
    fn from(e: E) -> Self { ImageError(e.to_string()) }
}

pub type Result<T> = std::result::Result<T, ImageError>;

/// The widest and the tallest picture a reader must accept (SPEC §2, *Reader limits*):
/// 65 536 pixels, the largest side an AV1 sequence header can declare.
pub const IMAGE_MAX_SIDE: u32 = 1 << 16;

/// The most pixels a reader must accept in one picture (SPEC §2, *Reader limits*):
/// 2^27, a 16 384 × 8 192 picture. Decoded to 8-bit RGB it is 384 MiB.
pub const IMAGE_MAX_PIXELS: u64 = 1 << 27;

/// Refuses a picture past the reader limits of SPEC §2 (`image_too_large`), given the
/// width and height its payload declares. The writers refuse the same pictures: a
/// conforming encoder writes no payload a conforming reader may refuse.
pub fn check_dimensions(w: u64, h: u64) -> Result<()> {
    if w > IMAGE_MAX_SIDE as u64 || h > IMAGE_MAX_SIDE as u64 || w * h > IMAGE_MAX_PIXELS {
        return Err(ImageError(format!(
            "image_too_large: the payload declares a {w} × {h} picture; a reader need accept at most {IMAGE_MAX_SIDE} on a side and {IMAGE_MAX_PIXELS} pixels (SPEC §2)"
        )));
    }
    Ok(())
}

/// Decodes a PNG or a JPEG payload, refusing any other format and any picture past
/// the reader limits, which it reads from the header before it decodes.
fn decode_as(bytes: &[u8], format: image::ImageFormat) -> Result<(Vec<u8>, u32, u32)> {
    let (w, h) = image::ImageReader::with_format(Cursor::new(bytes), format).into_dimensions()?;
    check_dimensions(w as u64, h as u64)?;
    let mut reader = image::ImageReader::with_format(Cursor::new(bytes), format);
    // The decoder's own buffer is at most 16-bit RGBA, 8 bytes a pixel; the default
    // allocation limit (512 MiB) would refuse such a picture within the reader limits.
    let mut limits = image::Limits::default();
    limits.max_image_width = Some(IMAGE_MAX_SIDE);
    limits.max_image_height = Some(IMAGE_MAX_SIDE);
    limits.max_alloc = Some(IMAGE_MAX_PIXELS * 8 + (1 << 24));
    reader.limits(limits);
    let img = reader.decode()?.to_rgb8();
    let (w, h) = img.dimensions();
    Ok((img.into_raw(), h, w))
}

// ---- PNG (lossless, universal) ----
pub fn png_encode(rgb: &[u8], h: u32, w: u32) -> Result<Vec<u8>> {
    check_dimensions(w as u64, h as u64)?;
    let mut out = Vec::new();
    image::codecs::png::PngEncoder::new(&mut out)
        .write_image(rgb, w, h, image::ExtendedColorType::Rgb8)?;
    Ok(out)
}

pub fn png_decode(bytes: &[u8]) -> Result<(Vec<u8>, u32, u32)> {
    decode_as(bytes, image::ImageFormat::Png)
}

// ---- JPEG (lossy) ----
// Uses `jpeg-encoder` (pure Rust) directly. Notes:
//   * `image` crate's default JpegEncoder produced files ~60-90% larger
//     than libjpeg at the same nominal q (bench caught it).
//   * `jpeg-encoder`'s q scale is offset from libjpeg's: q=80 there ≈
//     libjpeg q=92 in RD position. We remap the *public* WAI q to a
//     jpeg-encoder q so WAI q=N gives roughly libjpeg q=N output bytes,
//     letting users reason about WAI q the same way they reason about
//     libjpeg q (the universal convention in the field).
//   * We DON'T enable optimized Huffman tables — they're standard JPEG
//     but the `zune-jpeg` decoder we link via `image` misreads them.
#[cfg(feature = "codecs")]
fn remap_q_to_jpeg_encoder(public_q: u8) -> u8 {
    // Empirically calibrated against libjpeg on Kodak (bench harness):
    //   user-q 50 → encoder-q 35 ≈ libjpeg q=50 bytes
    //   user-q 80 → encoder-q 65 ≈ libjpeg q=80 bytes
    //   user-q 90 → encoder-q 78 ≈ libjpeg q=90 bytes
    // Linear-ish in the working range; clamp to 1..=100.
    let q = public_q.clamp(1, 100) as i32;
    let mapped = if q >= 50 { ((q - 50) * 130 / 100) + 35 } else { q * 70 / 100 };
    mapped.clamp(1, 100) as u8
}

#[cfg(feature = "codecs")]
pub fn jpeg_encode(rgb: &[u8], h: u32, w: u32, quality: u8) -> Result<Vec<u8>> {
    check_dimensions(w as u64, h as u64)?;
    let mut out = Vec::new();
    let q = remap_q_to_jpeg_encoder(quality);
    let enc = jpeg_encoder::Encoder::new(&mut out, q);
    enc.encode(rgb, w as u16, h as u16,
               jpeg_encoder::ColorType::Rgb)
        .map_err(|e| ImageError(format!("{e:?}")))?;
    Ok(out)
}

#[cfg(feature = "codecs")]
pub fn jpeg_decode(bytes: &[u8]) -> Result<(Vec<u8>, u32, u32)> {
    decode_as(bytes, image::ImageFormat::Jpeg)
}

// ---- AVIF (modern lossy, AV1-based) ----
#[cfg(feature = "codecs")]
pub fn avif_encode(rgb: &[u8], h: u32, w: u32, quality: u8) -> Result<Vec<u8>> {
    check_dimensions(w as u64, h as u64)?;
    let mut out = Vec::new();
    image::codecs::avif::AvifEncoder::new_with_speed_quality(&mut out, 6, quality)
        .write_image(rgb, w, h, image::ExtendedColorType::Rgb8)?;
    Ok(out)
}

#[cfg(feature = "codecs")]
pub fn avif_decode(bytes: &[u8]) -> Result<(Vec<u8>, u32, u32)> {
    // libavif directly. The `image` crate's own AVIF decoder is encode-only in the
    // default build; relying on it silently returned empty buffers (bench harness
    // caught this). libavif checks the size every image item declares (its `ispe`)
    // against the decoder's limits while it parses the container, before it decodes,
    // and checks the AV1 sequence header against them again, so the limits are set to
    // the reader limits rather than libavif's defaults (whose 32 768 side would refuse
    // a picture SPEC §2 says a reader must accept).
    use libavif_sys as sys;
    struct Decoder(*mut sys::avifDecoder);
    impl Drop for Decoder {
        fn drop(&mut self) { unsafe { sys::avifDecoderDestroy(self.0) } }
    }
    struct Image(*mut sys::avifImage);
    impl Drop for Image {
        fn drop(&mut self) { unsafe { sys::avifImageDestroy(self.0) } }
    }
    // SAFETY: each pointer libavif returns is checked for null before use and freed by
    // its guard; `bytes` outlives the decode that reads it; `out` is sized for the RGB
    // image libavif writes (`rowBytes` × `height`) before libavif writes it.
    unsafe {
        let dec = Decoder(sys::avifDecoderCreate());
        let img = Image(sys::avifImageCreateEmpty());
        if dec.0.is_null() || img.0.is_null() {
            return Err(ImageError("libavif decode: out of memory".into()));
        }
        (*dec.0).imageSizeLimit = IMAGE_MAX_PIXELS as u32;
        (*dec.0).imageDimensionLimit = IMAGE_MAX_SIDE;
        let r = sys::avifDecoderReadMemory(dec.0, img.0, bytes.as_ptr(), bytes.len());
        if r != sys::AVIF_RESULT_OK {
            let diag = std::ffi::CStr::from_ptr((*dec.0).diag.error.as_ptr()).to_string_lossy();
            if diag.contains("too large") {
                return Err(ImageError(format!(
                    "image_too_large: {diag}; a reader need accept at most {IMAGE_MAX_SIDE} on a side and {IMAGE_MAX_PIXELS} pixels (SPEC §2)"
                )));
            }
            return Err(ImageError(format!("libavif decode: result {r}: {diag}")));
        }
        let (w, h) = ((*img.0).width, (*img.0).height);
        check_dimensions(w as u64, h as u64)?;
        let mut rgb = sys::avifRGBImage::default();
        sys::avifRGBImageSetDefaults(&mut rgb, img.0);
        rgb.format = sys::AVIF_RGB_FORMAT_RGB;
        rgb.depth = 8;
        rgb.rowBytes = w * 3;
        let mut out = vec![0u8; w as usize * h as usize * 3];
        rgb.pixels = out.as_mut_ptr();
        let r = sys::avifImageYUVToRGB(img.0, &mut rgb);
        if r != sys::AVIF_RESULT_OK {
            return Err(ImageError(format!("libavif YUV to RGB: result {r}")));
        }
        Ok((out, h, w))
    }
}

// ---- JPEG-XL (lossless or lossy) ----
#[cfg(feature = "codecs")]
pub fn jxl_encode(rgb: &[u8], h: u32, w: u32, quality: Option<f32>) -> Result<Vec<u8>> {
    check_dimensions(w as u64, h as u64)?;
    // `quality=None` ⇒ mathematically lossless (the `.lossless(true)`
    // flag is what enables this; distance=0 alone is *not* sufficient).
    // `quality=Some(q)` ⇒ lossy with butteraugli distance derived from q.
    let mut enc = if let Some(q) = quality {
        let distance = ((100.0 - q.clamp(1.0, 100.0)) * 0.15).max(0.01);
        encoder_builder().speed(EncoderSpeed::Squirrel).quality(distance).build()?
    } else {
        // True lossless needs both `lossless(true)` AND
        // `uses_original_profile(true)` so libjxl keeps the original
        // color encoding instead of round-tripping through XYB internally
        // (which causes the small per-pixel drift we were seeing).
        encoder_builder()
            .speed(EncoderSpeed::Squirrel)
            .lossless(true)
            .quality(0.0)
            .uses_original_profile(true)
            .build()?
    };
    let r = enc.encode::<u8, u8>(rgb, w, h)?;
    Ok(r.data)
}

#[cfg(feature = "codecs")]
/// The width and height a JPEG XL payload declares, read by libjxl from its header
/// alone (the decoder is asked for nothing past the basic information).
fn jxl_dimensions(bytes: &[u8]) -> Result<(u32, u32)> {
    use jpegxl_sys::decode::{
        JxlDecoderCloseInput, JxlDecoderCreate, JxlDecoderDestroy, JxlDecoderGetBasicInfo,
        JxlDecoderProcessInput, JxlDecoderSetInput, JxlDecoderStatus, JxlDecoderSubscribeEvents,
    };
    struct Decoder(*mut jpegxl_sys::decode::JxlDecoder);
    impl Drop for Decoder {
        fn drop(&mut self) { unsafe { JxlDecoderDestroy(self.0) } }
    }
    // SAFETY: the decoder is checked for null and freed by its guard; `bytes` outlives
    // every call that reads it (input is closed and the decoder destroyed before this
    // returns); `info` is written by libjxl only when it reports success.
    unsafe {
        let dec = Decoder(JxlDecoderCreate(std::ptr::null()));
        if dec.0.is_null() {
            return Err(ImageError("jxl: out of memory".into()));
        }
        if JxlDecoderSubscribeEvents(dec.0, JxlDecoderStatus::BasicInfo as i32) != JxlDecoderStatus::Success
            || JxlDecoderSetInput(dec.0, bytes.as_ptr(), bytes.len()) != JxlDecoderStatus::Success
        {
            return Err(ImageError("jxl: decoder setup failed".into()));
        }
        JxlDecoderCloseInput(dec.0);
        if JxlDecoderProcessInput(dec.0) != JxlDecoderStatus::BasicInfo {
            return Err(ImageError("jxl: no header to read".into()));
        }
        let mut info = std::mem::MaybeUninit::<jpegxl_sys::metadata::codestream_header::JxlBasicInfo>::zeroed();
        if JxlDecoderGetBasicInfo(dec.0, info.as_mut_ptr()) != JxlDecoderStatus::Success {
            return Err(ImageError("jxl: no header to read".into()));
        }
        let info = info.assume_init();
        Ok((info.xsize, info.ysize))
    }
}

#[cfg(feature = "codecs")]
pub fn jxl_decode(bytes: &[u8]) -> Result<(Vec<u8>, u32, u32)> {
    let (w, h) = jxl_dimensions(bytes)?;
    check_dimensions(w as u64, h as u64)?;
    // Three 8-bit channels whatever the source holds: libjxl converts a grey picture to
    // RGB, drops alpha and rounds a deeper one to 8 bits, so every JPEG XL payload
    // decodes to the RGB this module promises.
    let dec = decoder_builder()
        .pixel_format(jpegxl_rs::decode::PixelFormat { num_channels: 3, ..Default::default() })
        .build()?;
    let (info, pixels) = dec.decode_with::<u8>(bytes)?;
    check_dimensions(info.width as u64, info.height as u64)?;
    if pixels.len() != info.width as usize * info.height as usize * 3 {
        return Err(ImageError(format!(
            "jxl: {} bytes decoded for a {} × {} picture", pixels.len(), info.width, info.height
        )));
    }
    Ok((pixels, info.height, info.width))
}

// The tests exercise every codec here: a build with `codecs` runs them; the floor's are
// tests/floor.rs.
#[cfg(all(test, feature = "codecs"))]
mod tests {
    use super::*;

    fn synth(h: u32, w: u32) -> Vec<u8> {
        let mut rgb = vec![0u8; (h * w * 3) as usize];
        for i in 0..h {
            for j in 0..w {
                let idx = ((i * w + j) * 3) as usize;
                rgb[idx]     = ((i + j) % 256) as u8;
                rgb[idx + 1] = ((i * 3) % 256) as u8;
                rgb[idx + 2] = ((j * 5) % 256) as u8;
            }
        }
        rgb
    }

    #[test]
    fn png_round_trip_lossless() {
        let (h, w) = (64u32, 64u32);
        let rgb = synth(h, w);
        let bytes = png_encode(&rgb, h, w).unwrap();
        let (rec, rh, rw) = png_decode(&bytes).unwrap();
        assert_eq!((rh, rw), (h, w));
        assert_eq!(rec, rgb, "PNG must be lossless");
    }

    #[test]
    fn jpeg_round_trip_lossy_but_sane() {
        let (h, w) = (64u32, 64u32);
        let rgb = synth(h, w);
        let bytes = jpeg_encode(&rgb, h, w, 90).unwrap();
        let (rec, _, _) = jpeg_decode(&bytes).unwrap();
        assert_eq!(rec.len(), rgb.len());
    }

    #[test]
    fn avif_round_trip_lossy_but_sane() {
        let (h, w) = (64u32, 64u32);
        let rgb = synth(h, w);
        let bytes = avif_encode(&rgb, h, w, 80).unwrap();
        assert!(bytes.len() > 0);
        // decode requires AVIF feature in `image` — may not be enabled
        // depending on build features; allow either case.
        let _ = avif_decode(&bytes);
    }

    #[test]
    fn jxl_lossless_round_trip() {
        let (h, w) = (64u32, 64u32);
        let rgb = synth(h, w);
        let bytes = jxl_encode(&rgb, h, w, None).unwrap();
        let (rec, rh, rw) = jxl_decode(&bytes).unwrap();
        assert_eq!((rh, rw), (h, w));
        assert_eq!(rec, rgb, "JPEG-XL distance=0 must be lossless");
    }
}
