inline_images.rs

//! Promote base64 / data-URL image payloads embedded in user text into real
//! multimodal attachments.
//!
//! Over SSH / headless TUI sessions, clipboard image paste often arrives as a
//! raw base64 blob or a `data:image/...;base64,...` URL rather than a file path.
//! The TUI already stages those at paste time when it can, but any path that
//! still lands base64 in the prompt text (paste missed, SDK/web send, replay)
//! must be recognized here so the model actually *sees* the image.

use base64::Engine;

/// Per-image byte cap after decode (matches TUI `maxAttachImageBytes`).
const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024;
/// Cap how many inline images we promote from one prompt (DoS / token blow-up).
const MAX_INLINE_IMAGES: usize = 8;
/// Minimum base64 compact length before we attempt a pure-base64 decode.
/// A real 1×1 PNG is ~68 raw / ~92 b64 chars; 64 is a floor that still rejects
/// short identifiers while accepting tiny screenshots.
const MIN_PURE_B64_LEN: usize = 64;

/// Result of scanning a user prompt for embedded image payloads.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct InlineImageExtract {
    /// Prompt text with data-URL / pure-base64 image blobs stripped.
    pub text: String,
    /// Data URLs (`data:image/<type>;base64,...`) ready for `image_to_data_url`
    /// / multimodal ContentPart construction. Order preserved; de-duplicated.
    pub images: Vec<String>,
}

/// Scan `text` for:
/// 1. `data:image/<subtype>;base64,<payload>` tokens (possibly mixed with prose)
/// 2. A whole-string pure base64 blob that decodes to PNG/JPEG/GIF/WEBP/BMP
///
/// Returns residual text + promoted data URLs. Non-image content is left
/// untouched. Fail-open: anything that does not clearly decode as an image
/// stays in the text so code pastes are never eaten.
pub(crate) fn extract_inline_images(text: &str) -> InlineImageExtract {
    if text.is_empty() {
        return InlineImageExtract {
            text: String::new(),
            images: Vec::new(),
        };
    }

    let mut images: Vec<String> = Vec::new();
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();

    // 1) data:image/...;base64,... URLs — strip them out of residual text.
    let (urls, rest) = extract_data_url_images(text);
    for url in urls {
        if images.len() >= MAX_INLINE_IMAGES {
            break;
        }
        if seen.insert(url.clone()) {
            images.push(url);
        }
    }

    let mut rest = rest;

    // 2) Whole residual (or whole original when no data URLs) is pure base64 of
    // an image — common for clipboard bridges over SSH.
    if images.len() < MAX_INLINE_IMAGES {
        let candidate = if rest.trim().is_empty() && images.is_empty() {
            text
        } else {
            rest.as_str()
        };
        if let Some(url) = try_decode_pure_base64_image(candidate) {
            if seen.insert(url.clone()) {
                images.push(url);
            }
            // Pure-base64 paste is fully consumed.
            return InlineImageExtract {
                text: String::new(),
                images,
            };
        }
    }

    // 3) Residual still has long base64 *tokens* mixed with prose ("what is
    // this?\n<b64...>") — promote each token that decodes to a real image and
    // strip it from the text. Fail-open for non-image tokens.
    if images.len() < MAX_INLINE_IMAGES && !rest.trim().is_empty() {
        let (more, cleaned) = extract_base64_tokens(&rest, MAX_INLINE_IMAGES - images.len());
        for url in more {
            if seen.insert(url.clone()) {
                images.push(url);
            }
        }
        rest = cleaned;
    }

    InlineImageExtract {
        text: rest.trim().to_string(),
        images,
    }
}

/// Pull every `data:image/...;base64,...` token out of `text`.
fn extract_data_url_images(text: &str) -> (Vec<String>, String) {
    const PREFIX: &str = "data:image/";
    let mut urls = Vec::new();
    let mut out = String::with_capacity(text.len());
    let bytes = text.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        // Find next data:image/ occurrence from i.
        let tail = &text[i..];
        if let Some(rel) = tail.find(PREFIX) {
            let start = i + rel;
            // Copy text before the match.
            out.push_str(&text[i..start]);
            // Scan the data URL body until whitespace / quote / end.
            let mut j = start + PREFIX.len();
            while j < bytes.len() {
                let c = bytes[j];
                if c == b' ' || c == b'\t' || c == b'\n' || c == b'\r' || c == b'"' || c == b'\'' {
                    break;
                }
                j += 1;
            }
            let candidate = &text[start..j];
            if let Some(url) = validate_data_url_image(candidate) {
                if urls.len() < MAX_INLINE_IMAGES {
                    urls.push(url);
                }
                // Skip the URL in residual text.
                i = j;
                continue;
            }
            // Not a valid image data URL — keep the prefix char and advance one
            // so we don't re-match the same spot forever.
            out.push(text[start..].chars().next().unwrap_or('d'));
            i = start + 1;
        } else {
            out.push_str(tail);
            break;
        }
    }
    (urls, out)
}

/// Validate a candidate data URL: must be `data:image/<type>;base64,<b64>` and
/// decode to a recognized image (or at least a non-trivial payload). Returns a
/// normalized data URL on success.
fn validate_data_url_image(candidate: &str) -> Option<String> {
    if !candidate.starts_with("data:image/") || !candidate.contains(";base64,") {
        return None;
    }
    let comma = candidate.find(',')?;
    if comma + 1 >= candidate.len() {
        return None;
    }
    let header = &candidate[..comma];
    let payload = &candidate[comma + 1..];
    // media type: data:image/<subtype>;base64
    let media = header
        .strip_prefix("data:")?
        .split(';')
        .next()
        .unwrap_or("image/png");
    if !media.starts_with("image/") {
        return None;
    }
    let raw = decode_base64(payload)?;
    if raw.len() > MAX_IMAGE_BYTES {
        return None;
    }
    // Prefer magic-byte confirmation; allow long payloads without magic only
    // when the declared type is a known image/* (provider will still reject junk).
    let sniffed = sniff_image_mime(&raw);
    if sniffed.is_none() && raw.len() <= 32 {
        return None;
    }
    let mime = sniffed.unwrap_or(media);
    // Re-encode with standard base64 so providers get a clean payload (input
    // may have had whitespace stripped already via the token scan).
    let b64 = base64::engine::general_purpose::STANDARD.encode(&raw);
    Some(format!("data:{mime};base64,{b64}"))
}

/// Scan residual text for long base64 tokens (≥ MIN_PURE_B64_LEN) that decode
/// to a real image. Returns promoted data URLs + residual with those tokens
/// removed. Adjacent pure-base64 lines are joined (clipboard line-wrap) before
/// decode so a multi-line dump mixed with a short question still works.
fn extract_base64_tokens(text: &str, budget: usize) -> (Vec<String>, String) {
    if budget == 0 || text.is_empty() {
        return (Vec::new(), text.to_string());
    }
    let lines: Vec<&str> = text.split('\n').collect();
    let mut images = Vec::new();
    let mut out_lines: Vec<String> = Vec::new();
    let mut i = 0;
    while i < lines.len() {
        let line = lines[i];
        let trimmed = line.trim();
        // Start of a potential base64 run: long-ish alphabet-only line.
        if is_base64_line(trimmed) && trimmed.len() >= 40 {
            let start = i;
            let mut run = trimmed.to_string();
            i += 1;
            while i < lines.len() {
                let t = lines[i].trim();
                if is_base64_line(t) && !t.is_empty() {
                    run.push_str(t);
                    i += 1;
                } else {
                    break;
                }
            }
            // Also accept a single long token that was space-separated on one line.
            if run.len() >= MIN_PURE_B64_LEN {
                if let Some(url) = try_decode_pure_base64_image(&run) {
                    if images.len() < budget {
                        images.push(url);
                        // Drop the consumed lines (don't push to out_lines).
                        continue;
                    }
                }
            }
            // Not an image — keep original lines.
            for keep in &lines[start..i] {
                out_lines.push((*keep).to_string());
            }
            continue;
        }
        // Per-whitespace-token scan on mixed lines ("see <b64> please").
        let mut rebuilt = String::new();
        let mut changed = false;
        for (ti, tok) in line.split_whitespace().enumerate() {
            if images.len() < budget && tok.len() >= MIN_PURE_B64_LEN {
                if let Some(url) = try_decode_pure_base64_image(tok) {
                    images.push(url);
                    changed = true;
                    continue;
                }
            }
            if ti > 0 || !rebuilt.is_empty() {
                // Preserve a single space between kept tokens.
                if !rebuilt.is_empty() {
                    rebuilt.push(' ');
                }
            }
            rebuilt.push_str(tok);
        }
        if changed {
            if !rebuilt.is_empty() {
                out_lines.push(rebuilt);
            }
        } else {
            out_lines.push(line.to_string());
        }
        i += 1;
    }
    let cleaned = out_lines.join("\n");
    (images, cleaned)
}

fn is_base64_line(s: &str) -> bool {
    !s.is_empty()
        && s.chars().all(|c| {
            c.is_ascii_alphanumeric() || c == '+' || c == '/' || c == '=' || c == '-' || c == '_'
        })
}

/// Treat the whole string as base64 of an image. Returns a data URL on success.
fn try_decode_pure_base64_image(s: &str) -> Option<String> {
    let s = s.trim();
    if s.starts_with("data:") {
        return None;
    }
    if s.len() < MIN_PURE_B64_LEN || s.len() > MAX_IMAGE_BYTES.saturating_mul(2) {
        return None;
    }
    // Collapse whitespace (clipboard bridges sometimes wrap lines).
    let compact: String = s.chars().filter(|c| !c.is_whitespace()).collect();
    if compact.len() < MIN_PURE_B64_LEN {
        return None;
    }
    // Reject anything that isn't base64 alphabet (plus URL-safe variant).
    if !compact.chars().all(|c| {
        c.is_ascii_alphanumeric() || c == '+' || c == '/' || c == '=' || c == '-' || c == '_'
    }) {
        return None;
    }
    // Heuristic: pure base64 of an image is almost entirely one token. If the
    // original (pre-collapse) has many mixed word tokens, leave it alone —
    // that looks like prose/code, not an image dump.
    let word_tokens = s.split_whitespace().filter(|t| !t.is_empty()).count();
    if word_tokens > 4 {
        // Allow a few wrapped lines of base64; more than ~4 whitespace-separated
        // chunks of non-b64-looking content is already filtered by alphabet.
        // Real wrapped base64 can be many lines though — only bail when tokens
        // look short (code identifiers) rather than long b64 lines.
        let short_tokens = s
            .split_whitespace()
            .filter(|t| !t.is_empty() && t.len() < 40)
            .count();
        if short_tokens > 2 {
            return None;
        }
    }
    let raw = decode_base64(&compact)?;
    if raw.len() > MAX_IMAGE_BYTES {
        return None;
    }
    let mime = sniff_image_mime(&raw)?;
    let b64 = base64::engine::general_purpose::STANDARD.encode(&raw);
    Some(format!("data:{mime};base64,{b64}"))
}

fn decode_base64(s: &str) -> Option<Vec<u8>> {
    let eng = base64::engine::general_purpose::STANDARD;
    if let Ok(v) = eng.decode(s) {
        return Some(v);
    }
    // Raw (no padding).
    let raw_eng = base64::engine::general_purpose::STANDARD_NO_PAD;
    if let Ok(v) = raw_eng.decode(s) {
        return Some(v);
    }
    // URL-safe.
    let url_eng = base64::engine::general_purpose::URL_SAFE;
    if let Ok(v) = url_eng.decode(s) {
        return Some(v);
    }
    let url_raw = base64::engine::general_purpose::URL_SAFE_NO_PAD;
    url_raw.decode(s).ok()
}

/// Sniff image MIME from magic bytes. Returns None for non-images.
pub(crate) fn sniff_image_mime(bytes: &[u8]) -> Option<&'static str> {
    if bytes.starts_with(&[0x89, b'P', b'N', b'G']) {
        Some("image/png")
    } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
        Some("image/jpeg")
    } else if bytes.starts_with(b"GIF8") {
        Some("image/gif")
    } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
        Some("image/webp")
    } else if bytes.starts_with(&[0x42, 0x4D]) {
        Some("image/bmp")
    } else if bytes.starts_with(b"<svg") || bytes.starts_with(b"<?xml") {
        // SVG is text; only accept when it looks like svg markup.
        let head = std::str::from_utf8(&bytes[..bytes.len().min(256)]).ok()?;
        if head.contains("<svg") {
            Some("image/svg+xml")
        } else {
            None
        }
    } else {
        None
    }
}

/// Merge explicit `images` (paths / data URLs from the client) with any
/// payloads promoted from prompt text. De-duplicates while preserving order
/// (explicit first, then inline).
pub(crate) fn merge_image_refs(explicit: Option<Vec<String>>, inline: Vec<String>) -> Vec<String> {
    let mut out = Vec::new();
    let mut seen = std::collections::HashSet::new();
    for img in explicit.into_iter().flatten().chain(inline) {
        let t = img.trim();
        if t.is_empty() {
            continue;
        }
        if seen.insert(t.to_string()) {
            out.push(t.to_string());
        }
        if out.len() >= MAX_INLINE_IMAGES {
            break;
        }
    }
    out
}

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

    /// Minimal valid 1×1 red PNG (69 bytes).
    fn tiny_png() -> Vec<u8> {
        let b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC";
        base64::engine::general_purpose::STANDARD
            .decode(b64)
            .expect("fixture")
    }

    fn tiny_png_b64() -> String {
        base64::engine::general_purpose::STANDARD.encode(tiny_png())
    }

    #[test]
    fn pure_base64_png_is_promoted() {
        let b64 = tiny_png_b64();
        let got = extract_inline_images(&b64);
        assert!(
            got.text.is_empty(),
            "residual should be empty, got {:?}",
            got.text
        );
        assert_eq!(got.images.len(), 1);
        assert!(got.images[0].starts_with("data:image/png;base64,"));
        // Round-trip: payload decodes back to PNG magic.
        let payload = got.images[0].split(',').nth(1).unwrap();
        let raw = base64::engine::general_purpose::STANDARD
            .decode(payload)
            .unwrap();
        assert_eq!(sniff_image_mime(&raw), Some("image/png"));
    }

    #[test]
    fn data_url_mixed_with_prose() {
        let b64 = tiny_png_b64();
        let url = format!("data:image/png;base64,{b64}");
        let text = format!("please look at {url} and describe it");
        let got = extract_inline_images(&text);
        assert_eq!(got.images.len(), 1);
        assert!(got.text.contains("please look at"));
        assert!(got.text.contains("and describe it"));
        assert!(!got.text.contains("data:image"));
    }

    #[test]
    fn pure_base64_mixed_with_prose() {
        // SSH/TUI case: short question + pasted raw base64 blob.
        let b64 = tiny_png_b64();
        let text = format!("what is this?\n{b64}");
        let got = extract_inline_images(&text);
        assert_eq!(
            got.images.len(),
            1,
            "mixed prose+b64 must promote the image"
        );
        assert!(got.images[0].starts_with("data:image/png;base64,"));
        assert!(got.text.contains("what is this"));
        assert!(
            !got.text.contains(&b64),
            "b64 blob must be stripped from residual"
        );
    }

    #[test]
    fn pure_base64_same_line_as_prose() {
        let b64 = tiny_png_b64();
        let text = format!("describe {b64} please");
        let got = extract_inline_images(&text);
        assert_eq!(got.images.len(), 1);
        assert!(got.text.contains("describe"));
        assert!(got.text.contains("please"));
        assert!(!got.text.contains(&b64));
    }

    #[test]
    fn plain_prose_untouched() {
        let text = "fix world\nsecond line of code: let x = 1;";
        let got = extract_inline_images(text);
        assert!(got.images.is_empty());
        assert_eq!(got.text.trim(), text);
    }

    #[test]
    fn short_base64_not_promoted() {
        // Too short to be a real image — must not eat identifiers.
        let got = extract_inline_images("YWJjZGVmZ2hpams="); // "abcdefghijk"
        assert!(got.images.is_empty());
    }

    #[test]
    fn code_looking_base64_alphabet_rejected_without_magic() {
        // Long base64-alphabet string that does NOT decode to an image.
        let s = "A".repeat(200);
        let got = extract_inline_images(&s);
        assert!(got.images.is_empty(), "non-image must stay text");
    }

    #[test]
    fn wrapped_base64_lines_accepted() {
        let b64 = tiny_png_b64();
        // Split into 40-char lines like some clipboard bridges do.
        let mut wrapped = String::new();
        for (i, c) in b64.chars().enumerate() {
            if i > 0 && i % 40 == 0 {
                wrapped.push('\n');
            }
            wrapped.push(c);
        }
        let got = extract_inline_images(&wrapped);
        assert_eq!(got.images.len(), 1, "wrapped b64 should promote");
        assert!(got.text.is_empty());
    }

    #[test]
    fn jpeg_magic_promoted() {
        // Minimal JPEG SOI + a few bytes (not a full image, but magic matches).
        // We need enough base64 length; pad with zeros after SOI/APP-ish bytes.
        let mut raw = vec![0xFF, 0xD8, 0xFF, 0xE0];
        raw.extend(std::iter::repeat(0u8).take(80));
        let b64 = base64::engine::general_purpose::STANDARD.encode(&raw);
        let got = extract_inline_images(&b64);
        assert_eq!(got.images.len(), 1);
        assert!(got.images[0].starts_with("data:image/jpeg;base64,"));
    }

    #[test]
    fn merge_explicit_then_inline() {
        let merged = merge_image_refs(
            Some(vec![
                "/tmp/a.png".into(),
                "data:image/png;base64,AA==".into(),
            ]),
            vec![
                "data:image/png;base64,AA==".into(), // dup
                "data:image/png;base64,BB==".into(),
            ],
        );
        assert_eq!(
            merged,
            vec![
                "/tmp/a.png".to_string(),
                "data:image/png;base64,AA==".to_string(),
                "data:image/png;base64,BB==".to_string(),
            ]
        );
    }

    #[test]
    fn empty_input() {
        let got = extract_inline_images("");
        assert!(got.text.is_empty() && got.images.is_empty());
    }
}