//! Images of chat messages are stored in the `blobs` table instead of as base64 inside the
//! message JSON. A chat then loads without its images; the UI fetches them on demand from
//! /api/blobs/<hash> (cached by the browser), and the runner puts the bytes back in before
//! a model request. Large PNGs without transparency are stored as JPEG.

use crate::types::{ImageData, Message, Part};
use base64::Engine;
use sha2::{Digest, Sha256};

const B64: base64::engine::GeneralPurpose = base64::engine::general_purpose::STANDARD;

/// PNGs above this size are re-encoded as JPEG when that is smaller and loses no transparency.
const RECOMPRESS_ABOVE: usize = 100 * 1024;
const JPEG_QUALITY: u8 = 85;

/// An image to put into the store.
pub struct Blob {
    pub hash: String,
    pub media_type: String,
    pub bytes: Vec<u8>,
}

/// Call `f` for every image of a message (pasted images and tool results).
pub fn for_each_image(m: &mut Message, mut f: impl FnMut(&mut ImageData)) {
    for part in &mut m.parts {
        match part {
            Part::Image { image } => f(image),
            Part::ToolResult { images, .. } => images.iter_mut().for_each(&mut f),
            _ => {}
        }
    }
}

/// Move the inline images of `m` out: each becomes a reference (`blob`, `size`, no
/// `data`). Returns the blobs to store. Images that already are references stay as they are.
pub fn externalize(m: &mut Message) -> Vec<Blob> {
    let mut out = Vec::new();
    for_each_image(m, |img| {
        if img.data.is_empty() {
            return;
        }
        let Ok(raw) = B64.decode(img.data.as_bytes()) else { return };
        let (media_type, bytes) = compact(&img.media_type, raw);
        let hash = hash(&bytes);
        img.media_type = media_type.clone();
        img.data.clear();
        img.blob = Some(hash.clone());
        img.size = Some(bytes.len() as u64);
        out.push(Blob { hash, media_type, bytes });
    });
    out
}

/// Whether `s` looks like a blob hash (64 lowercase hex digits).
pub fn is_hash(s: &str) -> bool {
    s.len() == 64 && s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
}

/// Address of `bytes` in the image store (SHA-256, hex).
pub fn hash(bytes: &[u8]) -> String {
    Sha256::digest(bytes).iter().map(|b| format!("{b:02x}")).collect()
}

/// Re-encode a large opaque PNG as JPEG if that is smaller; anything else is kept.
fn compact(media_type: &str, raw: Vec<u8>) -> (String, Vec<u8>) {
    if media_type != "image/png" || raw.len() <= RECOMPRESS_ABOVE {
        return (media_type.to_string(), raw);
    }
    let Ok(img) = image::load_from_memory_with_format(&raw, image::ImageFormat::Png) else {
        return (media_type.to_string(), raw);
    };
    if img.color().has_alpha() && img.to_rgba8().pixels().any(|p| p.0[3] != 255) {
        return (media_type.to_string(), raw);
    }
    let mut jpeg = Vec::new();
    let enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg, JPEG_QUALITY);
    if img.to_rgb8().write_with_encoder(enc).is_err() || jpeg.len() >= raw.len() {
        return (media_type.to_string(), raw);
    }
    ("image/jpeg".to_string(), jpeg)
}

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

    fn png(w: u32, h: u32, alpha: u8) -> Vec<u8> {
        // Noise, so the PNG is large and JPEG wins.
        let img = image::RgbaImage::from_fn(w, h, |x, y| {
            let v = (x.wrapping_mul(7919) ^ y.wrapping_mul(104_729)).wrapping_mul(2_654_435_761) >> 24;
            image::Rgba([v as u8, (v >> 1) as u8, (v >> 2) as u8, alpha])
        });
        let mut buf = Vec::new();
        img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png).unwrap();
        buf
    }

    fn message(img: ImageData) -> Message {
        let mut m = Message { id: "m1".into(), role: Role::User, parts: vec![], model: None, usage: None, created_at: 0 };
        m.parts.push(Part::ToolResult { call_id: "c".into(), name: "browser".into(), content: String::new(), images: vec![img], is_error: false, meta: None });
        m
    }

    #[test]
    fn large_opaque_png_becomes_a_jpeg_reference() {
        let raw = png(400, 400, 255);
        let mut m = message(ImageData::new("image/png", B64.encode(&raw)));
        let blobs = externalize(&mut m);
        assert_eq!(blobs.len(), 1);
        assert_eq!(blobs[0].media_type, "image/jpeg");
        assert!(blobs[0].bytes.len() < raw.len());
        let Part::ToolResult { images, .. } = &m.parts[0] else { panic!() };
        assert!(images[0].data.is_empty());
        assert_eq!(images[0].blob.as_deref(), Some(blobs[0].hash.as_str()));
        assert_eq!(images[0].size, Some(blobs[0].bytes.len() as u64));
        assert!(is_hash(&blobs[0].hash));
        // Already a reference: nothing to do.
        assert!(externalize(&mut m).is_empty());
    }

    #[test]
    fn transparent_and_small_pngs_stay_png() {
        let mut m = message(ImageData::new("image/png", B64.encode(png(400, 400, 128))));
        assert_eq!(externalize(&mut m)[0].media_type, "image/png");
        let mut m = message(ImageData::new("image/png", B64.encode(png(8, 8, 255))));
        assert_eq!(externalize(&mut m)[0].media_type, "image/png");
    }

    #[test]
    fn serialized_reference_has_no_data() {
        let img = ImageData { media_type: "image/jpeg".into(), data: String::new(), blob: Some("ab".into()), size: Some(3) };
        assert_eq!(serde_json::to_value(&img).unwrap(), serde_json::json!({"media_type": "image/jpeg", "blob": "ab", "size": 3}));
    }
}