//! DeepSeek V4 Flash Vision-Exp image layout and prefill masks.
//! Behavioral reference: DS4 0aaea5a238fb41a35106a551e73c8409dfb751ac.
use image::{DynamicImage, ImageDecoder, ImageReader};
use sha2::{Digest, Sha256};
use std::io::Cursor;

pub(super) const SOURCE_REVISION: &[u8] = b"e46e16bf6035c6f317eb2ac7458eb0362926d402";
pub(super) const START: i32 = 0;
pub(super) const PAD: i32 = 1;
pub(super) const IMAGE: i32 = 2;
pub(super) const NEWLINE: i32 = 3;
pub(super) const END: i32 = 4;

pub(super) struct Patches {
    pub values: Vec<f32>,
    pub fingerprint: [u8; 32],
    pub width: u32,
    pub height: u32,
    pub content_width: u32,
    pub content_height: u32,
    pub grid_width: u32,
    pub grid_height: u32,
}

pub(super) fn decode(encoded: &[u8]) -> Result<Patches, String> {
    if encoded.is_empty() || encoded.len() > 64 * 1024 * 1024 {
        return Err("image is empty or exceeds the 64 MiB encoded limit".into());
    }
    let reader = ImageReader::new(Cursor::new(encoded))
        .with_guessed_format()
        .map_err(|e| e.to_string())?;
    let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?;
    let (width, height) = decoder.dimensions();
    if width == 0
        || height == 0
        || width > 16384
        || height > 16384
        || u64::from(width) * u64::from(height) > 64 * 1024 * 1024
    {
        return Err("invalid DeepSeek image dimensions".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 rgb = image.into_rgb8();
    preprocess(rgb.as_raw(), rgb.width(), rgb.height())
}

fn grid_tokens(height: u32, width: u32) -> u32 {
    let h = (height / 14).div_ceil(3);
    let w = (width / 14).div_ceil(3);
    (h + h % 2) * (w + 1) + 2 + (h.div_ceil(2) * (w + 1) % 2) * 2
}

fn resize_shape(height: u32, width: u32) -> Result<(u32, u32), String> {
    let mut pw = width.min(height * 8);
    let mut ph = height;
    if u64::from(pw) * u64::from(ph) < 147456 {
        let scale = (147456.0 / (f64::from(pw) * f64::from(ph))).sqrt();
        pw = ((f64::from(pw) * scale) as u32).max(1);
        ph = ((f64::from(ph) * scale) as u32).max(1);
    }
    let mut bw = pw.div_ceil(14) * 14;
    let mut bh = ph.div_ceil(14) * 14;
    let mut budget = 381;
    while grid_tokens(bh, bw) > 381 {
        if budget <= 4 {
            return Err("image cannot fit the DeepSeek vision token budget".into());
        }
        let ratio = f64::from(ph) / f64::from(pw);
        let mw = ((f64::from(budget) - 2.0) / ratio + 0.25).sqrt() - 0.5;
        let mh = mw * ratio;
        if mw < 1.0 {
            bw = 42;
            bh = (((budget - 2) / 2) & !1) * 42;
        } else if mh < 2.0 {
            bh = 84;
            bw = ((budget - 2) / 2 - 1) * 42;
            if bw <= 42 {
                return Err("image cannot fit the DeepSeek vision token budget".into());
            }
        } else {
            let mw = mw.floor() as u32;
            let mh = mh.floor() as u32 & !1;
            let scale =
                (f64::from(mw) * 42.0 / f64::from(pw)).min(f64::from(mh) * 42.0 / f64::from(ph));
            bw = (f64::from(pw) * scale / 14.0).floor() as u32 * 14;
            bh = (f64::from(ph) * scale / 14.0).floor() as u32 * 14;
        }
        if bw == 0 || bh == 0 {
            return Err("invalid DeepSeek resized image dimensions".into());
        }
        budget -= 1;
    }
    Ok((bh, bw))
}

pub(super) fn preprocess(rgb: &[u8], width: u32, height: u32) -> Result<Patches, String> {
    if width == 0
        || height == 0
        || width > 16384
        || height > 16384
        || u64::from(width) * u64::from(height) > 64 * 1024 * 1024
        || rgb.len() != width as usize * height as usize * 3
    {
        return Err("invalid DeepSeek image dimensions".into());
    }
    let (bh, bw) = resize_shape(height, width)?;
    if bw > 16384 || bh > 16384 || u64::from(bw) * u64::from(bh) > 64 * 1024 * 1024 {
        return Err("resized DeepSeek image exceeds the decoded image limits".into());
    }
    let (cw, ch, ox, oy) = if width >= height * 8 {
        (bw, bh, 0, 0)
    } else {
        let scale = (f64::from(bw) / f64::from(width)).min(f64::from(bh) / f64::from(height));
        let cw = (f64::from(width) * scale)
            .round_ties_even()
            .clamp(1.0, f64::from(bw)) as u32;
        let ch = (f64::from(height) * scale)
            .round_ties_even()
            .clamp(1.0, f64::from(bh)) as u32;
        (
            cw,
            ch,
            (f64::from(bw - cw) * 0.5).round_ties_even() as u32,
            (f64::from(bh - ch) * 0.5).round_ties_even() as u32,
        )
    };
    let mut canvas = vec![127.0_f32; bw as usize * bh as usize * 3];
    let offset = (oy as usize * bw as usize + ox as usize) * 3;
    super::metal::vision::resize_bicubic(rgb, width, height, &mut canvas[offset..], cw, ch, bw);
    for value in &mut canvas {
        *value = *value / 127.5 - 1.0;
    }
    let gw = bw / 14;
    let gh = bh / 14;
    let mut values = Vec::with_capacity(gw as usize * gh as usize * 588);
    for py in 0..gh {
        for px in 0..gw {
            for c in 0..3 {
                for y in 0..14 {
                    for x in 0..14 {
                        values.push(
                            canvas[((py * 14 + y) as usize * bw as usize + (px * 14 + x) as usize)
                                * 3
                                + c],
                        );
                    }
                }
            }
        }
    }
    let mut hash = Sha256::new();
    hash.update(width.to_le_bytes());
    hash.update(height.to_le_bytes());
    hash.update(rgb);
    Ok(Patches {
        fingerprint: hash.finalize().into(),
        values,
        width,
        height,
        content_width: cw,
        content_height: ch,
        grid_width: gw,
        grid_height: gh,
    })
}

pub(super) struct Layout {
    pub types: Vec<i32>,
    pub permutation: Vec<usize>,
}

pub(super) fn layout(height: u32, width: u32, position: u32) -> Result<Layout, String> {
    if height == 0 || width == 0 {
        return Err("invalid DeepSeek image token grid".into());
    }
    let rows = u64::from(height) + u64::from(height % 2);
    let row_length = u64::from(width) + 1;
    let leading = 3 - position % 4;
    let trailing = ((rows / 2 * row_length) % 2) * 2;
    let count = u64::from(leading) + 2 + rows * row_length + trailing;
    if count > 384 {
        return Err("DeepSeek image token layout exceeds its model limit".into());
    }
    let mut types = vec![PAD; leading as usize];
    let mut permutation = Vec::with_capacity(height as usize * width as usize);
    types.push(START);
    for pair in 0..rows / 2 {
        for column in 0..row_length {
            for row in pair * 2..pair * 2 + 2 {
                if row < u64::from(height) && column < u64::from(width) {
                    types.push(IMAGE);
                    permutation.push((row * u64::from(width) + column) as usize);
                } else {
                    types.push(if row < u64::from(height) {
                        NEWLINE
                    } else {
                        PAD
                    });
                }
            }
        }
    }
    types.extend(std::iter::repeat_n(PAD, trailing as usize));
    types.push(END);
    Ok(Layout { types, permutation })
}

#[derive(Debug, PartialEq)]
pub(super) struct Span {
    pub block: usize,
    pub start: usize,
    pub end: usize,
}

pub(super) fn spans(tokens: &[i32], vocab: i32) -> Result<Vec<Span>, String> {
    let mut result = Vec::new();
    let mut cursor = 0;
    while cursor < tokens.len() {
        if (0..vocab).contains(&tokens[cursor]) {
            cursor += 1;
            continue;
        }
        let block = cursor;
        while cursor < tokens.len() && tokens[cursor] == vocab + PAD {
            cursor += 1;
        }
        if tokens.get(cursor) != Some(&(vocab + START)) {
            return Err("malformed DeepSeek image block".into());
        }
        let start = cursor;
        cursor += 1;
        loop {
            let token = *tokens
                .get(cursor)
                .ok_or("unterminated DeepSeek image block")?;
            if token <= vocab + START || token > vocab + END {
                return Err("malformed DeepSeek image block".into());
            }
            if token == vocab + END {
                break;
            }
            cursor += 1;
        }
        result.push(Span {
            block,
            start,
            end: cursor,
        });
        cursor += 1;
    }
    Ok(result)
}

pub(super) fn prefill_chunk(tokens: &[i32], vocab: i32, cap: u32) -> Result<usize, String> {
    let mut limit = tokens.len().min(cap as usize);
    if limit == 0 {
        return Err("empty DeepSeek prefill chunk".into());
    }
    for span in spans(tokens, vocab)? {
        if span.block >= limit {
            break;
        }
        if span.end + 1 > limit {
            if span.end + 1 > cap as usize {
                if span.block == 0 {
                    return Err("DeepSeek image block exceeds prefill capacity".into());
                }
                limit = span.block;
                break;
            }
            limit = span.end + 1;
        }
    }
    Ok(limit)
}

pub(super) fn attention_bounds(
    tokens: &[i32],
    vocab: i32,
    pos: u32,
    n_raw: u32,
    window: u32,
) -> Result<Vec<u32>, String> {
    let end = pos
        .checked_add(
            tokens
                .len()
                .try_into()
                .map_err(|_| "vision attention length overflow")?,
        )
        .ok_or("vision attention position overflow")?;
    if tokens.is_empty() || n_raw == 0 || n_raw > end {
        return Err("invalid DeepSeek vision attention extent".into());
    }
    let first = end - n_raw;
    let mut bounds = Vec::with_capacity(tokens.len() * 2);
    for q in 0..tokens.len() as u32 {
        let p = pos + q;
        bounds.extend([
            first.max(if window == 0 {
                0
            } else {
                (p + 1).saturating_sub(window)
            }),
            p.min(end - 1),
        ]);
    }
    for span in spans(tokens, vocab)? {
        for q in span.start..=span.end {
            let p = pos + q as u32;
            let left = (q - span.start).min(383) as u32;
            let right = (span.end - q).min(384) as u32;
            let add = if window != 0 && left >= window {
                left - (window - 1)
            } else {
                0
            };
            bounds[2 * q] = first.max(if window == 0 {
                0
            } else {
                p.saturating_sub(window - 1 + add)
            });
            bounds[2 * q + 1] = (p + right).min(end - 1);
        }
    }
    Ok(bounds)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn issue100_ds4_image_layout_and_attention_vectors() {
        for (h, w, p, types, perm) in [
            (
                2,
                3,
                0,
                vec![1, 1, 1, 0, 2, 2, 2, 2, 2, 2, 3, 3, 4],
                vec![0, 3, 1, 4, 2, 5],
            ),
            (
                3,
                2,
                5,
                vec![1, 1, 0, 2, 2, 2, 2, 3, 3, 2, 1, 2, 1, 3, 1, 4],
                vec![0, 2, 1, 3, 4, 5],
            ),
            (1, 1, 3, vec![0, 2, 1, 3, 1, 4], vec![0]),
        ] {
            let actual = layout(h, w, p).unwrap();
            assert_eq!(actual.types, types);
            assert_eq!(actual.permutation, perm);
        }
        let tokens = [7, 101, 101, 100, 102, 103, 104, 8, 100, 101, 104, 9];
        assert_eq!(
            spans(&tokens, 100).unwrap(),
            vec![
                Span {
                    block: 1,
                    start: 3,
                    end: 6
                },
                Span {
                    block: 8,
                    start: 8,
                    end: 10
                }
            ]
        );
        assert_eq!(prefill_chunk(&tokens, 100, 5).unwrap(), 1);
        assert_eq!(prefill_chunk(&tokens[1..], 100, 6).unwrap(), 6);
        assert!(prefill_chunk(&tokens[1..], 100, 5).is_err());
        assert!(prefill_chunk(&tokens[4..], 100, 8).is_err());
        for malformed in [
            [101, 4, 5, 6],
            [100, 4, 104, 6],
            [100, 100, 104, 6],
            [100, 102, 6, 7],
            [100, 105, 6, 7],
            [-1, 100, 104, 6],
        ] {
            assert!(spans(&malformed, 100).is_err());
        }
        assert_eq!(
            attention_bounds(
                &[7, 101, 100, 102, 102, 102, 102, 102, 104, 8],
                100,
                10,
                20,
                4
            )
            .unwrap(),
            vec![
                7, 10, 8, 11, 9, 18, 10, 18, 11, 18, 12, 18, 12, 18, 12, 18, 12, 18, 16, 19
            ]
        );
    }
    #[test]
    fn issue100_ds4_preprocess_dimensions_and_limits() {
        let mut rgb = Vec::new();
        for y in 0..9 {
            for x in 0..17 {
                rgb.extend([
                    (x * 13 + y * 3) as u8,
                    (x * 5 + y * 17) as u8,
                    (x * 7 + y * 11) as u8,
                ]);
            }
        }
        let p = preprocess(&rgb, 17, 9).unwrap();
        assert_eq!(
            (
                p.grid_width,
                p.grid_height,
                p.content_width,
                p.content_height
            ),
            (38, 20, 529, 280)
        );
        assert_eq!(p.values.len(), 760 * 588);
        assert!(
            p.values
                .iter()
                .all(|v| v.is_finite() && (-1.0..=1.0).contains(v))
        );
        for (h, w) in [
            (1, 1),
            (1, 8),
            (1, 16384),
            (16384, 1),
            (384, 384),
            (1024, 1024),
            (1080, 1920),
            (4096, 4096),
        ] {
            let (bh, bw) = resize_shape(h, w).unwrap();
            assert!(grid_tokens(bh, bw) <= 381);
            for start in 0..4 {
                assert!(
                    layout((bh / 14).div_ceil(3), (bw / 14).div_ceil(3), start)
                        .unwrap()
                        .types
                        .len()
                        <= 384
                );
            }
        }
        assert!(preprocess(&[], 0, 1).is_err());
        assert!(layout(384, 384, 0).is_err());
    }
}