//! dllm-vision: CPU image decode + Gemma-4 image preprocessing.
//!
//! Produces exactly what the HF `Gemma4ImageProcessor` (torchvision backend, the one `AutoProcessor`
//! selects for this model, also used by vLLM) feeds the vision tower, as `u8` patches:
//!
//!  1. decode + PIL-style `convert('RGB')`            (`decode`)
//!  2. `get_aspect_ratio_preserving_size`              (this file)
//!  3. torchvision uint8 antialiased bicubic resize    (`resize`, bit-exact with aten, NOT with PIL.resize)
//!  4. patchify (C,gh,16,gw,16)->(gh,gw,16,16,C)       (fused into the vertical resize pass)
//!
//! Rescale: HF computes `float32(x_u8) * float32(1/255)` (rescale_factor = 0.00392156862745098 as
//! float32 multiply), which is NOT always equal to `x/255.0f`; the GPU side must use
//! `x as f32 * 0.003921569f32` (= f32 of 1/255) to be bit-identical, then `2*(x-0.5)` happens in
//! the model.  Pixels are therefore kept as `u8` here.

pub mod decode;
pub mod pool;
pub mod png_fast;
pub mod resize;

pub use decode::{decode_to_rgb, probe, Format};

use rayon::prelude::*;
use resize::Rgb;

/// PIL raises DecompressionBombError above 2 * 89_478_485 pixels.
pub const DEFAULT_MAX_PIXELS: usize = 2 * 89_478_485;

#[derive(Debug)]
pub enum VisionError {
    Decode(String),
    TooLarge { w: usize, h: usize, max_pixels: usize },
    Unsupported,
    /// get_aspect_ratio_preserving_size ValueError analogues
    Geometry(String),
}

impl std::fmt::Display for VisionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VisionError::Decode(s) => write!(f, "cannot decode image: {s}"),
            VisionError::TooLarge { w, h, max_pixels } => write!(f, "image {w}x{h} exceeds the {max_pixels} pixel limit"),
            VisionError::Unsupported => write!(f, "unsupported image format"),
            VisionError::Geometry(s) => write!(f, "{s}"),
        }
    }
}
impl std::error::Error for VisionError {}
pub type Result<T> = std::result::Result<T, VisionError>;

#[derive(Clone, Debug)]
pub struct PreprocConfig {
    pub patch_size: usize,
    pub pooling_kernel_size: usize,
    pub max_soft_tokens: usize,
    /// decompression-bomb guard (decoded pixels)
    pub max_pixels: usize,
}

impl Default for PreprocConfig {
    /// the model's processor_config (bicubic, antialias, rescale 1/255, no normalize) at the default 280-token budget
    fn default() -> Self {
        PreprocConfig { patch_size: 16, pooling_kernel_size: 3, max_soft_tokens: 280, max_pixels: DEFAULT_MAX_PIXELS }
    }
}

#[derive(Clone, Debug)]
pub struct ImagePatches {
    /// resized size (multiples of pooling*patch = 48)
    pub height: usize,
    pub width: usize,
    /// patch grid (height/16, width/16)
    pub gh: usize,
    pub gw: usize,
    /// `[gh*gw, 768]` u8; patch index = gy*gw+gx, within a patch values ordered (dy, dx, c)
    pub patches: Vec<u8>,
    /// `gh*gw / pooling^2` (<= max_soft_tokens)
    pub n_soft: usize,
    /// 64-bit content hash of dims + patches (prefix-cache identity)
    pub hash: u64,
}

/// Port of `get_aspect_ratio_preserving_size` (HF gemma4), including its edge cases.
pub fn aspect_ratio_preserving_size(
    height: usize,
    width: usize,
    patch_size: usize,
    max_patches: usize,
    pooling: usize,
) -> Result<(usize, usize)> {
    let total_px = (height * width) as f64;
    let target_px = (max_patches * patch_size * patch_size) as f64;
    let factor = (target_px / total_px).sqrt();
    let ideal_h = factor * height as f64;
    let ideal_w = factor * width as f64;
    let side_mult = pooling * patch_size;
    let mut th = (ideal_h / side_mult as f64).floor() as i64 * side_mult as i64;
    let mut tw = (ideal_w / side_mult as f64).floor() as i64 * side_mult as i64;
    if th == 0 && tw == 0 {
        return Err(VisionError::Geometry(format!(
            "Attempting to resize to a 0 x 0 image. Resized height should be divisible by `pooling_kernel_size * patch_size`={side_mult}."
        )));
    }
    let max_side = ((max_patches / (pooling * pooling)) * side_mult) as i64;
    if th == 0 {
        th = side_mult as i64;
        tw = ((width as f64 / height as f64).floor() as i64 * side_mult as i64).min(max_side);
    } else if tw == 0 {
        tw = side_mult as i64;
        th = ((height as f64 / width as f64).floor() as i64 * side_mult as i64).min(max_side);
    }
    if (th * tw) as f64 > target_px {
        return Err(VisionError::Geometry(format!(
            "Resizing [{height}x{width}] to [{th}x{tw}] but this exceeds {max_patches} patches with patch_size {patch_size}"
        )));
    }
    Ok((th as usize, tw as usize))
}

impl PreprocConfig {
    pub fn target_size(&self, height: usize, width: usize) -> Result<(usize, usize)> {
        aspect_ratio_preserving_size(
            height,
            width,
            self.patch_size,
            self.max_soft_tokens * self.pooling_kernel_size * self.pooling_kernel_size,
            self.pooling_kernel_size,
        )
    }
}

/// n_soft (and resized size) from the image dimensions alone, e.g. for context-length estimation
/// from `probe()`.
pub fn n_soft_for_size(width: usize, height: usize, cfg: &PreprocConfig) -> Result<usize> {
    let (th, tw) = cfg.target_size(height, width)?;
    Ok((th / cfg.patch_size) * (tw / cfg.patch_size) / (cfg.pooling_kernel_size * cfg.pooling_kernel_size))
}

/// Decode + preprocess one image.
pub fn preprocess(bytes: &[u8], cfg: &PreprocConfig) -> Result<ImagePatches> {
    if decode::sniff(bytes) == Format::Png {
        if let Some(p) = png_fast::parse(bytes) {
            if png_fast::worthwhile(&p) {
                return preprocess_png_fast(&p, cfg);
            }
        }
    }
    let rgb = decode::decode_rgb(bytes, cfg.max_pixels)?;
    patches_from_rgb(&rgb, cfg)
}

/// PNG fast path: inflate+unfilter on a helper thread while this thread (and the rayon pool) runs
/// the horizontal resize pass on the rows that are already finished.
fn preprocess_png_fast(p: &png_fast::Parsed, cfg: &PreprocConfig) -> Result<ImagePatches> {
    use std::sync::atomic::{AtomicUsize, Ordering};
    let (w, h) = (p.w, p.h);
    if w.checked_mul(h).map_or(true, |px| px > cfg.max_pixels) {
        return Err(VisionError::TooLarge { w, h, max_pixels: cfg.max_pixels });
    }
    let (th, tw) = cfg.target_size(h, w)?;
    let mut src = Rgb::scratch(w, h);
    let progress = AtomicUsize::new(0);
    let streaming = tw != w && h >= 256 && rayon::current_num_threads() >= 2;
    if !streaming {
        png_fast::decode_into(p, src.data.as_mut_ptr(), src.data.len(), &progress)?;
        return patches_from_rgb(&src, cfg);
    }
    let coeffs = resize::compute_coeffs(w, tw);
    let mid = Rgb::scratch(tw, h);
    let sptr = resize::SendPtr(src.data.as_mut_ptr());
    let slen = src.data.len();
    let mid_ptr = resize::SendPtr(mid.data.as_ptr() as *mut u8);
    let min_batch = (h / 32).max(16);
    let res: Result<()> = std::thread::scope(|sc| {
        let sp = sptr.clone();
        let prog = &progress;
        let dec = sc.spawn(move || {
            let sp = sp;
            png_fast::decode_into(p, sp.0, slen, prog)
        });
        let mut hdone = 0usize;
        while hdone < h {
            let done = progress.load(Ordering::Acquire);
            // keep one row of lag: the SIMD kernel may read a few bytes past a row's end
            let avail = if done >= h { h } else { done.saturating_sub(1) };
            if avail >= hdone + min_batch || (avail == h && avail > hdone) {
                unsafe { resize::h_rows_seq(sptr.0 as *const u8, w * 3, &coeffs, hdone, avail, mid_ptr.0, tw * 3) };
                hdone = avail;
            } else if dec.is_finished() && progress.load(Ordering::Acquire) < h {
                break; // decoder failed
            } else {
                std::hint::spin_loop();
                std::thread::yield_now();
            }
        }
        match dec.join() {
            Ok(r) => r,
            Err(_) => Err(VisionError::Decode("PNG decoder thread panicked".into())),
        }
    });
    res?;
    drop(src);
    finish(&mid, th, tw, cfg)
}

/// Resize + patchify an already decoded image.
pub fn patches_from_rgb(src: &Rgb, cfg: &PreprocConfig) -> Result<ImagePatches> {
    let (th, tw) = cfg.target_size(src.h, src.w)?;
    if tw != src.w {
        let mid = resize::resize_h(src, tw);
        finish(&mid, th, tw, cfg)
    } else {
        finish(src, th, tw, cfg)
    }
}

/// `mid` is already `tw` wide (horizontal pass done / not needed); vertical pass (if `th != mid.h`)
/// fused with patchify, then the content hash.
fn finish(mid: &Rgb, th: usize, tw: usize, cfg: &PreprocConfig) -> Result<ImagePatches> {
    debug_assert_eq!(mid.w, tw);
    let ps = cfg.patch_size;
    let (gh, gw) = (th / ps, tw / ps);
    let pp = ps * ps * 3; // 768
    let row_bytes = ps * 3; // 48
    let mut patches = pool::take(gh * gw * pp); // fully overwritten below
    let vc = if th != mid.h { Some(resize::compute_coeffs(mid.h, th)) } else { None };
    let stride = tw * 3;
    patches.par_chunks_mut(gw * pp).enumerate().for_each_init(
        || vec![0u8; stride],
        |tmp, (gy, chunk)| {
            for dy in 0..ps {
                let y = gy * ps + dy;
                let row: &[u8] = match &vc {
                    Some(c) => {
                        resize::v_one(mid, c, y, tmp);
                        tmp
                    }
                    None => &mid.data[y * stride..(y + 1) * stride],
                };
                for gx in 0..gw {
                    chunk[gx * pp + dy * row_bytes..gx * pp + (dy + 1) * row_bytes]
                        .copy_from_slice(&row[gx * row_bytes..(gx + 1) * row_bytes]);
                }
            }
        },
    );
    let n_soft = gh * gw / (cfg.pooling_kernel_size * cfg.pooling_kernel_size);
    let mut hasher = xxhash_rust::xxh3::Xxh3::new();
    hasher.update(&(th as u64).to_le_bytes());
    hasher.update(&(tw as u64).to_le_bytes());
    hasher.update(&patches);
    let hash = hasher.digest();
    Ok(ImagePatches { height: th, width: tw, gh, gw, patches, n_soft, hash })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn geometry_matches_hf() {
        let c = PreprocConfig::default();
        // values taken from the real processor (see tests/golden)
        assert_eq!(c.target_size(353, 834).unwrap(), (480, 1200));
        assert_eq!(c.target_size(938, 1448).unwrap(), (624, 960));
        assert_eq!(c.target_size(1, 1000).unwrap(), (48, 13440));
        assert_eq!(c.target_size(1000, 1).unwrap(), (13440, 48));
        assert!(c.target_size(0 + 1, 1).is_ok());
    }
}