//! Image decoding to 8-bit RGB with exactly PIL's `Image.open(...).convert('RGB')` semantics.
//!
//!  * PNG: `png` crate (fdeflate).  Lossless, bit exact.  Alpha is dropped (not composited),
//!    palette/tRNS ignored for RGB conversion, 16-bit RGB(A)/LA keep the high byte, 16-bit
//!    gray ("I;16" in PIL) is clipped to 255 like PIL's `I16L_L`.
//!  * JPEG: libjpeg-turbo 3.x (vendored via `turbojpeg-sys`), the same library PIL uses: ISLOW IDCT,
//!    fancy upsampling, default libjpeg YCbCr->RGB.  CMYK/YCCK: PIL's inverted-Adobe convention
//!    and `cmyk2rgb`.
//!  * GIF/BMP/TIFF/WebP: `image` crate (best effort; lossy WebP is not guaranteed bit-exact
//!    with libwebp).

use crate::resize::{Rgb, PAD};
use crate::{Result, VisionError};
use std::io::Cursor;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Format {
    Png,
    Jpeg,
    Gif,
    Bmp,
    WebP,
    Tiff,
    Unknown,
}

pub fn sniff(b: &[u8]) -> Format {
    if b.len() >= 8 && b[..8] == [0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a] {
        Format::Png
    } else if b.len() >= 3 && b[0] == 0xff && b[1] == 0xd8 && b[2] == 0xff {
        Format::Jpeg
    } else if b.len() >= 6 && (&b[..6] == b"GIF87a" || &b[..6] == b"GIF89a") {
        Format::Gif
    } else if b.len() >= 2 && &b[..2] == b"BM" {
        Format::Bmp
    } else if b.len() >= 12 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" {
        Format::WebP
    } else if b.len() >= 4 && (b[..4] == [b'I', b'I', 42, 0] || b[..4] == [b'M', b'M', 0, 42]) {
        Format::Tiff
    } else {
        Format::Unknown
    }
}

fn err<T: std::fmt::Display>(e: T) -> VisionError {
    VisionError::Decode(e.to_string())
}

// ----------------------------------------------------------------------------- probe

/// Header-only dimension probe (width, height).
pub fn probe(bytes: &[u8]) -> Result<(usize, usize)> {
    match sniff(bytes) {
        Format::Png => {
            if bytes.len() < 24 || &bytes[12..16] != b"IHDR" {
                return Err(err("bad PNG header"));
            }
            let w = u32::from_be_bytes(bytes[16..20].try_into().unwrap()) as usize;
            let h = u32::from_be_bytes(bytes[20..24].try_into().unwrap()) as usize;
            Ok((w, h))
        }
        Format::Jpeg => probe_jpeg(bytes),
        Format::Gif => {
            if bytes.len() < 10 {
                return Err(err("bad GIF header"));
            }
            Ok((
                u16::from_le_bytes([bytes[6], bytes[7]]) as usize,
                u16::from_le_bytes([bytes[8], bytes[9]]) as usize,
            ))
        }
        Format::Unknown => Err(VisionError::Unsupported),
        _ => {
            let r = image::ImageReader::new(Cursor::new(bytes)).with_guessed_format().map_err(err)?;
            let (w, h) = r.into_dimensions().map_err(err)?;
            Ok((w as usize, h as usize))
        }
    }
}

fn probe_jpeg(b: &[u8]) -> Result<(usize, usize)> {
    let mut i = 2;
    while i + 4 <= b.len() {
        if b[i] != 0xff {
            i += 1;
            continue;
        }
        let m = b[i + 1];
        if m == 0xff {
            i += 1;
            continue;
        }
        if m == 0xd8 || (0xd0..=0xd7).contains(&m) || m == 0x01 || m == 0x00 {
            i += 2;
            continue;
        }
        let len = u16::from_be_bytes([b[i + 2], b[i + 3]]) as usize;
        // SOF0..SOF15 except DHT(c4), JPG(c8), DAC(cc)
        if (0xc0..=0xcf).contains(&m) && m != 0xc4 && m != 0xc8 && m != 0xcc {
            if i + 9 > b.len() {
                break;
            }
            let h = u16::from_be_bytes([b[i + 5], b[i + 6]]) as usize;
            let w = u16::from_be_bytes([b[i + 7], b[i + 8]]) as usize;
            return Ok((w, h));
        }
        if m == 0xda {
            break;
        }
        i += 2 + len;
    }
    Err(err("JPEG: no SOF marker"))
}

// ----------------------------------------------------------------------------- entry

/// Decode to packed RGB8 (PIL `convert('RGB')` semantics). `max_pixels` guards against
/// decompression bombs (checked from the header before any large allocation).
pub fn decode_rgb(bytes: &[u8], max_pixels: usize) -> Result<Rgb> {
    let fmt = sniff(bytes);
    if fmt == Format::Unknown {
        return Err(VisionError::Unsupported);
    }
    let (w, h) = probe(bytes)?;
    if w == 0 || h == 0 {
        return Err(err("zero-sized image"));
    }
    if w.checked_mul(h).map_or(true, |p| p > max_pixels) {
        return Err(VisionError::TooLarge { w, h, max_pixels });
    }
    match fmt {
        Format::Png => decode_png(bytes),
        Format::Jpeg => decode_jpeg(bytes),
        _ => decode_other(bytes),
    }
}

/// `(w, h, rgb)` with `rgb.len() == w*h*3`.
pub fn decode_to_rgb(bytes: &[u8]) -> Result<(usize, usize, Vec<u8>)> {
    let mut r = decode_rgb(bytes, crate::DEFAULT_MAX_PIXELS)?;
    let mut v = std::mem::take(&mut r.data);
    v.truncate(r.w * r.h * 3);
    Ok((r.w, r.h, v))
}

// ----------------------------------------------------------------------------- PNG

fn decode_png(bytes: &[u8]) -> Result<Rgb> {
    if let Some(p) = crate::png_fast::parse(bytes).filter(crate::png_fast::worthwhile) {
        let mut out = Rgb::scratch(p.w, p.h);
        let prog = std::sync::atomic::AtomicUsize::new(0);
        crate::png_fast::decode_into(&p, out.data.as_mut_ptr(), out.data.len(), &prog)?;
        return Ok(out);
    }
    decode_png_slow(bytes)
}

fn decode_png_slow(bytes: &[u8]) -> Result<Rgb> {
    use png::{BitDepth, ColorType, Transformations};
    let limits = png::Limits { bytes: 1usize << 40 };
    // first look at the header to special-case 16-bit gray (PIL "I;16")
    let probe_dec = png::Decoder::new_with_limits(Cursor::new(bytes), limits.clone());
    let probe_reader = probe_dec.read_info().map_err(err)?;
    let (w, h) = {
        let i = probe_reader.info();
        (i.width as usize, i.height as usize)
    };
    let gray16 = probe_reader.info().color_type == ColorType::Grayscale
        && probe_reader.info().bit_depth == BitDepth::Sixteen;
    drop(probe_reader);

    let mut dec = png::Decoder::new_with_limits(Cursor::new(bytes), limits);
    dec.set_transformations(if gray16 { Transformations::IDENTITY } else { Transformations::EXPAND | Transformations::STRIP_16 });
    let mut reader = dec.read_info().map_err(err)?;
    let (ct, bd) = reader.output_color_type();
    let size = reader.output_buffer_size().ok_or_else(|| err("PNG buffer size overflow"))?;
    // RGB output decodes straight into the final (padded) buffer
    let direct = ct == ColorType::Rgb && bd == BitDepth::Eight;
    let mut buf = vec![0u8; if direct { size + PAD } else { size }];
    reader.next_frame(&mut buf[..size]).map_err(err)?;
    if direct {
        return Ok(Rgb { w, h, data: buf });
    }
    let mut out = vec![0u8; w * h * 3 + PAD];
    let n = w * h;
    match (ct, bd) {
        (ColorType::Rgba, BitDepth::Eight) => rgba_to_rgb(&buf, &mut out[..n * 3]),
        (ColorType::Grayscale, BitDepth::Eight) => {
            for (i, &g) in buf[..n].iter().enumerate() {
                out[i * 3] = g;
                out[i * 3 + 1] = g;
                out[i * 3 + 2] = g;
            }
        }
        (ColorType::GrayscaleAlpha, BitDepth::Eight) => {
            for i in 0..n {
                let g = buf[i * 2];
                out[i * 3] = g;
                out[i * 3 + 1] = g;
                out[i * 3 + 2] = g;
            }
        }
        (ColorType::Grayscale, BitDepth::Sixteen) => {
            for i in 0..n {
                let (hi, lo) = (buf[i * 2], buf[i * 2 + 1]);
                let g = if hi != 0 { 255 } else { lo };
                out[i * 3] = g;
                out[i * 3 + 1] = g;
                out[i * 3 + 2] = g;
            }
        }
        _ => return Err(err(format!("unsupported PNG output {:?}/{:?}", ct, bd))),
    }
    Ok(Rgb { w, h, data: out })
}

pub fn rgba_to_rgb(src: &[u8], dst: &mut [u8]) {
    use rayon::prelude::*;
    dst.par_chunks_mut(3 * 4096).zip(src.par_chunks(4 * 4096)).for_each(|(d, s)| {
        for (dp, sp) in d.chunks_exact_mut(3).zip(s.chunks_exact(4)) {
            dp.copy_from_slice(&sp[..3]);
        }
    });
}

// ----------------------------------------------------------------------------- JPEG

fn decode_jpeg(bytes: &[u8]) -> Result<Rgb> {
    use turbojpeg_sys as tj;
    unsafe {
        let h = tj::tj3Init(tj::TJINIT_TJINIT_DECOMPRESS as i32);
        if h.is_null() {
            return Err(err("tj3Init failed"));
        }
        struct Guard(tj::tjhandle);
        impl Drop for Guard {
            fn drop(&mut self) {
                unsafe { tj::tj3Destroy(self.0) }
            }
        }
        let _g = Guard(h);
        let msg = |h: tj::tjhandle| -> VisionError {
            let p = tj::tj3GetErrorStr(h);
            let s = if p.is_null() { "jpeg error".to_string() } else { std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned() };
            VisionError::Decode(s)
        };
        // defaults == PIL: accurate DCT (ISLOW), fancy upsampling; keep going on warnings
        tj::tj3Set(h, tj::TJPARAM_TJPARAM_FASTDCT as i32, 0);
        tj::tj3Set(h, tj::TJPARAM_TJPARAM_FASTUPSAMPLE as i32, 0);
        if tj::tj3DecompressHeader(h, bytes.as_ptr(), bytes.len() as _) != 0 {
            return Err(msg(h));
        }
        let w = tj::tj3Get(h, tj::TJPARAM_TJPARAM_JPEGWIDTH as i32) as usize;
        let ht = tj::tj3Get(h, tj::TJPARAM_TJPARAM_JPEGHEIGHT as i32) as usize;
        let cs = tj::tj3Get(h, tj::TJPARAM_TJPARAM_COLORSPACE as i32);
        let prec = tj::tj3Get(h, tj::TJPARAM_TJPARAM_PRECISION as i32);
        if prec != 8 {
            return Err(err(format!("unsupported JPEG precision {prec}")));
        }
        if cs == tj::TJCS_TJCS_CMYK as i32 || cs == tj::TJCS_TJCS_YCCK as i32 {
            let mut cmyk = vec![0u8; w * ht * 4];
            if tj::tj3Decompress8(h, bytes.as_ptr(), bytes.len() as _, cmyk.as_mut_ptr(), (w * 4) as i32, tj::TJPF_TJPF_CMYK as i32) != 0 {
                return Err(msg(h));
            }
            let mut out = vec![0u8; w * ht * 3 + PAD];
            // PIL: rawmode "CMYK;I" (inverted, Adobe) then Convert.c cmyk2rgb
            for (o, p) in out.chunks_exact_mut(3).zip(cmyk.chunks_exact(4)) {
                let (c, m, y, k) = (255 - p[0] as u32, 255 - p[1] as u32, 255 - p[2] as u32, 255 - p[3] as u32);
                let nk = 255 - k;
                let md = |a: u32| -> u8 {
                    let t = a * nk + 128;
                    let v = ((t >> 8) + t) >> 8;
                    nk.saturating_sub(v).min(255) as u8
                };
                o[0] = md(c);
                o[1] = md(m);
                o[2] = md(y);
            }
            return Ok(Rgb { w, h: ht, data: out });
        }
        let mut out = vec![0u8; w * ht * 3 + PAD];
        if tj::tj3Decompress8(h, bytes.as_ptr(), bytes.len() as _, out.as_mut_ptr(), (w * 3) as i32, tj::TJPF_TJPF_RGB as i32) != 0 {
            return Err(msg(h));
        }
        Ok(Rgb { w, h: ht, data: out })
    }
}

// ----------------------------------------------------------------------------- others

fn decode_other(bytes: &[u8]) -> Result<Rgb> {
    let img = image::load_from_memory(bytes).map_err(err)?;
    let rgb = img.to_rgb8();
    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
    Ok(Rgb::from_vec(w, h, rgb.into_raw()))
}