//! PNG/JPEG preprocessing and image positions from MTPLX's Qwen3-VL path.
use image::{DynamicImage, ImageDecoder, ImageReader, RgbImage};
use std::io::Cursor;

pub(super) struct Image {
    pub(super) pixels: Vec<f32>,
    pub(super) h: i32,
    pub(super) w: i32,
}

pub(super) fn decode(bytes: &[u8]) -> Result<Image, String> {
    if bytes.is_empty() || bytes.len() > 50 * 1024 * 1024 {
        return Err("Qwen image exceeds the 50 MiB encoded limit".into());
    }
    let reader = ImageReader::new(Cursor::new(bytes))
        .with_guessed_format()
        .map_err(|e| e.to_string())?;
    if !matches!(
        reader.format(),
        Some(image::ImageFormat::Png | image::ImageFormat::Jpeg)
    ) {
        return Err("Qwen images must be PNG or JPEG".into());
    }
    let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?;
    let (w, h) = decoder.dimensions();
    if w == 0 || h == 0 || w > 8000 || h > 8000 {
        return Err("Qwen images must have dimensions between 1 and 8000".into());
    }
    let orientation = decoder.orientation().map_err(|e| e.to_string())?;
    let mut image = DynamicImage::from_decoder(decoder).map_err(|e| e.to_string())?;
    image.apply_orientation(orientation);
    let mut image = image.into_rgb8();
    let (w, h) = image.dimensions();
    let (rh, rw) = smart_resize(h, w)?;
    if (w, h) != (rw, rh) {
        image = resize(&image, rw, rh);
    }
    let (h, w) = ((rh / 16) as i32, (rw / 16) as i32);
    let mut pixels = Vec::with_capacity(h as usize * w as usize * 1536);
    for br in 0..h / 2 {
        for bc in 0..w / 2 {
            for ir in 0..2 {
                for ic in 0..2 {
                    for c in 0..3 {
                        for _t in 0..2 {
                            for y in 0..16 {
                                for x in 0..16 {
                                    let v = image.get_pixel(
                                        ((bc * 2 + ic) * 16 + x) as u32,
                                        ((br * 2 + ir) * 16 + y) as u32,
                                    )[c];
                                    pixels.push((f32::from(v) * (1_f32 / 255.) - 0.5) / 0.5);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    Ok(Image { pixels, h, w })
}
fn smart_resize(h: u32, w: u32) -> Result<(u32, u32), String> {
    if f64::from(h.max(w)) / f64::from(h.min(w)) > 200. {
        return Err("Qwen image aspect ratio exceeds 200".into());
    }
    let (hf, wf) = (f64::from(h), f64::from(w));
    let (mut rh, mut rw) = (
        (hf / 32.).round_ties_even() as u32 * 32,
        (wf / 32.).round_ties_even() as u32 * 32,
    );
    if u64::from(rh) * u64::from(rw) > 16777216 {
        let beta = (hf * wf / 16777216.).sqrt();
        rh = ((hf / beta / 32.).floor() as u32 * 32).max(32);
        rw = ((wf / beta / 32.).floor() as u32 * 32).max(32);
    } else if u64::from(rh) * u64::from(rw) < 65536 {
        let beta = (65536. / (hf * wf)).sqrt();
        rh = (hf * beta / 32.).ceil() as u32 * 32;
        rw = (wf * beta / 32.).ceil() as u32 * 32;
    }
    Ok((rh, rw))
}
// Pillow Resample.c: bicubic a=-0.5, normalized 22-bit coefficients,
// horizontal then vertical passes with rounding/clipping after each pass.
fn coefficients(input: u32, output: u32) -> Vec<(usize, Vec<i32>)> {
    let scale = f64::from(input) / f64::from(output);
    let filter_scale = scale.max(1.);
    let support = 2. * filter_scale;
    (0..output)
        .map(|i| {
            let center = (f64::from(i) + 0.5) * scale;
            let first = ((center - support + 0.5) as i64).max(0) as usize;
            let last = ((center + support + 0.5) as i64).min(i64::from(input)) as usize;
            let ws = (first..last)
                .map(|j| {
                    let x = ((j as f64 - center + 0.5) / filter_scale).abs();
                    if x < 1. {
                        ((1.5 * x - 2.5) * x) * x + 1.
                    } else if x < 2. {
                        (((x - 5.) * x + 8.) * x - 4.) * (-0.5)
                    } else {
                        0.
                    }
                })
                .collect::<Vec<_>>();
            let sum = ws.iter().sum::<f64>();
            (
                first,
                ws.iter()
                    .map(|w| (w / sum * 4194304.).round() as i32)
                    .collect(),
            )
        })
        .collect()
}
fn resize(input: &RgbImage, w: u32, h: u32) -> RgbImage {
    let horizontal = if w == input.width() {
        input.clone()
    } else {
        let coeff = coefficients(input.width(), w);
        RgbImage::from_fn(w, input.height(), |x, y| {
            let (start, ws) = &coeff[x as usize];
            image::Rgb(std::array::from_fn(|c| {
                let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| {
                    sum + i64::from(*k) * i64::from(input.get_pixel((*start + i) as u32, y)[c])
                });
                (sum >> 22).clamp(0, 255) as u8
            }))
        })
    };
    if h == horizontal.height() {
        return horizontal;
    }
    let coeff = coefficients(horizontal.height(), h);
    RgbImage::from_fn(w, h, |x, y| {
        let (start, ws) = &coeff[y as usize];
        image::Rgb(std::array::from_fn(|c| {
            let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| {
                sum + i64::from(*k) * i64::from(horizontal.get_pixel(x, (*start + i) as u32)[c])
            });
            (sum >> 22).clamp(0, 255) as u8
        }))
    })
}

pub(super) fn positions(ids: &[i32], grids: &[(i32, i32)]) -> Result<(Vec<i32>, i32), String> {
    let n = ids.len();
    let mut axes = vec![0; n * 3];
    let (mut start, mut next) = (0, 0_i32);
    for &(h, w) in grids {
        let h = h / 2;
        let w = w / 2;
        if h <= 0 || w <= 0 {
            return Err("invalid Qwen image grid".into());
        }
        let end = start
            + ids[start..]
                .iter()
                .position(|&t| t == 248056)
                .ok_or("missing Qwen image pad")?;
        for i in start..end {
            for axis in 0..3 {
                axes[axis * n + i] = next + (i - start) as i32;
            }
        }
        next += (end - start) as i32;
        let count = (h * w) as usize;
        if end + count > n || ids[end..end + count].iter().any(|&t| t != 248056) {
            return Err("Qwen image pad count differs from grid".into());
        }
        for i in 0..count {
            axes[end + i] = next;
            axes[n + end + i] = next + i as i32 / w;
            axes[n * 2 + end + i] = next + i as i32 % w;
        }
        next += h.max(w).max(1);
        start = end + count;
    }
    if ids[start..].contains(&248056) {
        return Err("Qwen image pads exceed image grids".into());
    }
    for i in start..n {
        for axis in 0..3 {
            axes[axis * n + i] = next + (i - start) as i32;
        }
    }
    next += (n - start) as i32;
    Ok((axes, next - n as i32))
}

#[test]
fn qwen_vision_positions_and_resize_contract() {
    assert_eq!(smart_resize(1024, 1024).unwrap(), (1024, 1024));
    assert_eq!(smart_resize(16, 16).unwrap(), (256, 256));
    assert!(smart_resize(1, 201).is_err());
    let (table, delta) = positions(&[1, 248056, 248056, 248056, 248056, 2], &[(4, 4)]).unwrap();
    assert_eq!(
        table,
        vec![0, 1, 1, 1, 1, 3, 0, 1, 1, 2, 2, 3, 0, 1, 2, 1, 2, 3]
    );
    assert_eq!(delta, -2);
    assert!(positions(&[248056], &[(4, 4)]).is_err());
}

#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(super) struct Identity {
    pub(super) start: usize,
    pub(super) end: usize,
    pub(super) digest: [u8; 32],
    /// Missing on checkpoints written before image-aware sparse QSA.
    #[serde(default)]
    pub(super) policy: u32,
}

pub(super) fn cache_policy() -> u32 {
    u32::from(std::env::var("MTPLX_QWEN4_VISION_QSA").as_deref() != Ok("0"))
}

pub(super) fn compatible(images: &[Identity]) -> bool {
    images.iter().all(|image| image.policy == cache_policy())
}

pub(super) fn same_prefix(a: &[Identity], b: &[Identity], tokens: usize) -> bool {
    // Image pads do not identify partial image embeddings. MTPLX only restores
    // before an image or after its complete span, never into the middle.
    !b.iter()
        .any(|image| image.start < tokens && tokens < image.end)
        && a.iter()
            .filter(|i| i.start < tokens)
            .eq(b.iter().filter(|i| i.start < tokens))
}

#[test]
fn qwen_image_preprocessing_matches_pillow_mtplx() {
    use sha2::{Digest, Sha256};
    for (w, h, gh, gw, expected) in [
        (
            17,
            31,
            22,
            12,
            "339e4ac4e02bb8fcab6d521194200191e4c77fbcce224de98b40d8d48393930f",
        ),
        (
            333,
            177,
            12,
            22,
            "3fc0890cb34e0c37bf32c5ce2a8bfa9b178ae3564c8e97b8e6496e71a66b2644",
        ),
        (
            1537,
            257,
            16,
            96,
            "96bcea68014063c4edc51b422af82930e1bd43984d57a41e12f44865eed45368",
        ),
    ] {
        let image = RgbImage::from_fn(w, h, |x, y| {
            image::Rgb(std::array::from_fn(|c| {
                ((x * 13 + y * 7 + c as u32 * 31) % 256) as u8
            }))
        });
        let mut bytes = Cursor::new(Vec::new());
        image.write_to(&mut bytes, image::ImageFormat::Png).unwrap();
        let input = decode(bytes.get_ref()).unwrap();
        assert_eq!((input.h, input.w), (gh, gw));
        let hash = Sha256::digest(
            input
                .pixels
                .iter()
                .flat_map(|v| v.to_le_bytes())
                .collect::<Vec<_>>(),
        );
        assert_eq!(
            hash.iter().map(|b| format!("{b:02x}")).collect::<String>(),
            expected,
            "{w}x{h}"
        );
    }
}

#[test]
fn qwen_image_cache_keys_distinguish_pixels_and_preserve_prior_text() {
    let old = vec![Identity {
        start: 10,
        end: 14,
        digest: [1; 32],
        policy: cache_policy(),
    }];
    let changed = vec![Identity {
        start: 10,
        end: 14,
        digest: [2; 32],
        policy: cache_policy(),
    }];
    assert!(same_prefix(&old, &old, 20));
    assert!(!same_prefix(&old, &changed, 20));
    assert!(same_prefix(&old, &changed, 10));
    assert!(!same_prefix(&old, &[], 20));
    assert!(!same_prefix(&old, &old, 11));
    assert!(!same_prefix(&old, &old, 13));
    assert!(same_prefix(&old, &old, 14));
}