choreo-daemon 0.2.0

use super::{PreparedImage, ToolExecError, context::ToolContext, truncate_tool_output};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use choreo_keystore::ServiceCredential;
use image::GenericImageView;
use resvg::usvg;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::{io, time::Duration};
use tracing::{debug, info, warn};
use url::Url;

#[derive(Debug, Deserialize, JsonSchema)]
pub struct DisplayImageArgs {
    /// MIME type of the image (e.g. "image/png", "image/svg+xml")
    mime_type: String,
    /// Path to an image file on disk
    path: Option<String>,
    /// URL of an image to fetch and display
    url: Option<String>,
    /// Base64-encoded image data
    base64_data: Option<String>,
    /// Raw SVG markup to render
    svg_text: Option<String>,
    /// Alt text description of the image
    alt: Option<String>,
}

pub(crate) const MAX_DISPLAY_IMAGE_BYTES: usize = 8 * 1024 * 1024;
const IMAGE_FETCH_TIMEOUT_SECS: u64 = 10;

/// The image MIME types `display_image` accepts. Covers every raster format the
/// `image` crate decodes, plus SVG (`resvg`), plus HEIC/HEIF (`heif-oxide`), and
/// AVIF (gated behind the `avif` feature — see `normalize_image_mime_type`).
/// The gate is intentionally broad: the client's decoder sniffs the bytes, and
/// `inspect_image_dimensions` validates it can actually decode the payload.
fn is_supported_image_mime(mime: &str) -> bool {
    matches!(
        mime,
        "image/png"
            | "image/jpeg"
            | "image/webp"
            | "image/gif"
            | "image/bmp"
            | "image/x-bmp"
            | "image/x-ms-bmp"
            | "image/tiff"
            | "image/tif"
            | "image/targa"
            | "image/x-tga"
            | "image/x-targa"
            | "image/vnd.microsoft.icon"
            | "image/x-icon"
            | "image/x-portable-anymap"
            | "image/x-portable-pixmap"
            | "image/x-portable-graymap"
            | "image/x-portable-bitmap"
            | "image/vnd.radiance"
            | "image/x-hdr"
            | "image/hdr"
            | "image/x-exr"
            | "image/openexr"
            | "image/qoi"
            | "image/x-dds"
            | "image/vnd.ms-dds"
            | "image/farbfeld"
            | "image/x-farbfeld"
            | "image/avif"
            | "image/heic"
            | "image/heif"
            | "image/svg+xml"
    )
}

/// The `display_image` tool's return value: a human-readable text handle plus
/// the prepared image, so the framework's `extract_image` hook reads the image
/// straight off the per-invocation return value (no shared state). `impl
/// Serialize` emits only `text`, so the JSON tool result is a plain string
/// exactly as before.
#[derive(Debug)]
pub struct DisplayImageReturn {
    /// The text handle (mime, dimensions, bytes) shown to the model.
    pub text: String,
    /// The prepared image handed to the client via `extract_image`.
    pub image: PreparedImage,
}

impl Serialize for DisplayImageReturn {
    /// Serialize to just the text handle, keeping the JSON wire format a plain
    /// string (identical to the previous `Return = String`).
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(&self.text)
    }
}

impl JsonSchema for DisplayImageReturn {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        std::borrow::Cow::Borrowed("DisplayImageReturn")
    }

    fn json_schema(_gen: &mut schemars::SchemaGenerator) -> schemars::Schema {
        schemars::json_schema!({ "type": "string" })
    }
}

fn prepare_image(args: &DisplayImageArgs) -> io::Result<PreparedImage> {
    // Normalized up front too (the shared helper re-normalizes idempotently)
    // so an unsupported MIME is rejected before any network/file I/O — the
    // original error precedence of `display_image`, preserved by the refactor.
    let mime_type = normalize_image_mime_type(&args.mime_type)?;
    let selected_sources = [
        args.path.as_ref().map(|_| "path"),
        args.url.as_ref().map(|_| "url"),
        args.base64_data.as_ref().map(|_| "base64_data"),
        args.svg_text.as_ref().map(|_| "svg_text"),
    ]
    .into_iter()
    .flatten()
    .count();
    if selected_sources != 1 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "provide exactly one image source: path, url, base64_data, or svg_text",
        ));
    }

    let data = if let Some(path) = &args.path {
        std::fs::read(path.trim())?
    } else if let Some(url) = &args.url {
        fetch_image_bytes(url.trim(), mime_type)?
    } else if let Some(base64_data) = &args.base64_data {
        BASE64.decode(base64_data.trim()).map_err(|error| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid base64_data: {error}"),
            )
        })?
    } else if let Some(svg_text) = &args.svg_text {
        svg_text.as_bytes().to_vec()
    } else {
        // Provably unreachable: `selected_sources != 1` above already rejects
        // every case other than exactly one Some source, so one of the four
        // branches took it. The daemon must never panic even on a logical
        // bug, so this defensive fallback returns a structured error instead
        // of unreachable! — a mismatch would otherwise crash the process.
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "no image source set",
        ));
    };

    // Shared with `generate_image` (tools/image_gen.rs): every image that
    // reaches a client display goes through the same normalize → cap →
    // dimension-probe pipeline, so a new image-producing tool cannot drift
    // from `display_image`'s safety posture (decompression-bomb guard, AVIF
    // feature gate, size ceiling).
    let (mime_type, width, height) = prepare_image_from_bytes(mime_type, &data)?;
    Ok(PreparedImage {
        mime_type,
        data,
        width,
        height,
        alt: args.alt.clone().filter(|alt| !alt.trim().is_empty()),
    })
}

/// Normalize a MIME type, enforce the display-size cap, and probe the pixel
/// dimensions of in-memory image bytes — the common tail of both image tools.
/// Returns the normalized mime plus `(width, height)`; `data` is returned to
/// the caller untouched so it can build its own [`PreparedImage`]. Kept
/// separate from source acquisition (path/url/base64/svg) because the
/// `generate_image` tool starts from provider bytes, not from user sources.
pub(crate) fn prepare_image_from_bytes(
    mime_type: &str,
    data: &[u8],
) -> io::Result<(String, u32, u32)> {
    let mime_type = normalize_image_mime_type(mime_type)?;

    if data.len() > MAX_DISPLAY_IMAGE_BYTES {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "image exceeds maximum allowed size of {}",
                humfmt::bytes(MAX_DISPLAY_IMAGE_BYTES as u64),
            ),
        ));
    }

    let (width, height) = inspect_image_dimensions(mime_type, data)?;
    Ok((mime_type.to_string(), width, height))
}

fn normalize_image_mime_type(mime_type: &str) -> io::Result<&str> {
    let normalized = mime_type.trim();
    if !is_supported_image_mime(normalized) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("unsupported image mime type: {normalized}"),
        ));
    }
    // AVIF is gated behind the `avif` feature (`image/avif-native`, a C
    // library). Recognized but rejected when the feature is off, keeping the
    // default/release build C-free.
    if normalized == "image/avif" && !cfg!(feature = "avif") {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "image/avif is gated behind the `avif` feature ".to_string(),
        ));
    }
    Ok(normalized)
}

fn fetch_image_bytes(url_str: &str, expected_mime_type: &str) -> io::Result<Vec<u8>> {
    let url =
        Url::parse(url_str).map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
    match url.scheme() {
        "http" | "https" => {}
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "image url must use http or https",
            ));
        }
    }

    let agent = ureq::Agent::new_with_config(
        ureq::Agent::config_builder()
            .timeout_global(Some(Duration::from_secs(IMAGE_FETCH_TIMEOUT_SECS)))
            .http_status_as_error(false)
            .build(),
    );
    let response = agent.get(url.as_str()).call().map_err(io::Error::other)?;
    let status = response.status().as_u16();
    if !(200..300).contains(&status) {
        return Err(io::Error::other(format!(
            "image request failed with status {status}"
        )));
    }
    if let Some(content_type) = response.headers().get("content-type")
        && let Ok(content_type) = content_type.to_str()
        && !content_type.starts_with(expected_mime_type)
    {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "image response content-type {content_type} does not match {expected_mime_type}"
            ),
        ));
    }
    let bytes = response
        .into_body()
        .read_to_vec()
        .map_err(io::Error::other)?;
    Ok(bytes)
}

fn inspect_image_dimensions(mime_type: &str, data: &[u8]) -> io::Result<(u32, u32)> {
    match mime_type {
        // All raster formats the `image` crate decodes (PNG, JPEG, WebP, GIF,
        // BMP, TIFF, TGA, DDS, ICO, PNM, HDR, OpenEXR, Farbfeld, QOI) — plus
        // AVIF when the gated `avif` feature is on. `decode_raster_oriented`
        // sniffs the bytes (so the MIME is largely advisory) *and* applies the
        // decompression-bomb `image::Limits` guard, so a hostile raster can't
        // drive a huge allocation during the dimension probe either — not just
        // the display decode.
        "image/png"
        | "image/jpeg"
        | "image/webp"
        | "image/gif"
        | "image/bmp"
        | "image/x-bmp"
        | "image/x-ms-bmp"
        | "image/tiff"
        | "image/tif"
        | "image/targa"
        | "image/x-tga"
        | "image/x-targa"
        | "image/vnd.microsoft.icon"
        | "image/x-icon"
        | "image/x-portable-anymap"
        | "image/x-portable-pixmap"
        | "image/x-portable-graymap"
        | "image/x-portable-bitmap"
        | "image/vnd.radiance"
        | "image/x-hdr"
        | "image/hdr"
        | "image/x-exr"
        | "image/openexr"
        | "image/qoi"
        | "image/x-dds"
        | "image/vnd.ms-dds"
        | "image/farbfeld"
        | "image/x-farbfeld"
        | "image/avif" => {
            // Route through the shared guarded decoder (with EXIF orientation
            // baked in), so a hostile raster is bounded before it allocates.
            let img = choreo_image::decode_raster_oriented(data).map_err(io::Error::other)?;
            Ok(img.dimensions())
        }
        "image/svg+xml" => {
            let options = usvg::Options::default();
            let tree = usvg::Tree::from_data(data, &options).map_err(io::Error::other)?;
            let size = tree.size().to_int_size();
            Ok((size.width(), size.height()))
        }
        "image/heic" | "image/heif" => {
            // Route through the shared guarded decoder so a hostile HEIC's
            // declared geometry can't drive a huge allocation during the
            // dimension probe either (not just the client display decode).
            let img = choreo_image::decode_heic(data).map_err(io::Error::other)?;
            Ok(img.dimensions())
        }
        _ => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("unsupported image mime type: {mime_type}"),
        )),
    }
}

pub(crate) struct DisplayImage {}

impl DisplayImage {
    pub(crate) fn new() -> Self {
        DisplayImage {}
    }
}

impl super::Tool for DisplayImage {
    type Args = DisplayImageArgs;
    type Return = DisplayImageReturn;
    type Error = ToolExecError;

    fn name(&self) -> &'static str {
        "display_image"
    }
    fn description(&self) -> &'static str {
        "Display an image (PNG, JPEG, WebP, GIF, BMP, TIFF, SVG, HEIC/HEIF, and more; AVIF behind the `avif` feature) in the client UI."
    }
    fn describe_invocation(&self, args: &Self::Args) -> String {
        let mut parts = vec![format!("Displaying image ({}).", args.mime_type)];
        if let Some(ref p) = args.path {
            parts.push(format!(" Path: `{}`.", p));
        }
        if let Some(ref u) = args.url {
            parts.push(format!(" URL: {}.", u));
        }
        if args.base64_data.is_some() {
            parts.push(" Source: base64 data.".to_string());
        }
        if args.svg_text.is_some() {
            parts.push(" Source: SVG markup.".to_string());
        }
        if let Some(ref alt) = args.alt {
            parts.push(format!(" Alt text: {}.", alt));
        }
        parts.concat()
    }

    fn return_string(ret: &Self::Return) -> String {
        ret.text.clone()
    }
    fn execute(
        &self,
        args: Self::Args,
        _x_credentials: Option<&ServiceCredential>,
        _working_dir: Option<&Path>,
        _ctx: Option<&ToolContext>,
    ) -> Result<Self::Return, Self::Error> {
        let image = match prepare_image(&args) {
            Ok(image) => image,
            Err(e) => {
                warn!(error = %e, "display_image: failed to prepare image");
                return Err(ToolExecError(e.to_string()));
            }
        };
        let mime_type = image.mime_type.clone();
        let width = image.width;
        let height = image.height;
        let byte_len = image.data.len();
        debug!(
            mime = %mime_type,
            width,
            height,
            bytes = byte_len,
            "display_image: prepared image"
        );
        let text = truncate_tool_output(&format!(
            "displayed image ({mime_type}, {width}x{height}, {})",
            humfmt::bytes(byte_len as u64),
        ));
        info!(
            mime = %mime_type,
            width,
            height,
            bytes = byte_len,
            "display_image: displayed image successfully"
        );
        // Carry the image in the return value for the framework's `extract_image`
        // hook to read from `ret` (no shared-state parking).
        Ok(DisplayImageReturn { text, image })
    }

    fn extract_image(&self, ret: &Self::Return) -> Option<PreparedImage> {
        Some(ret.image.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use image::ImageFormat;
    use std::io::Cursor;

    #[test]
    fn raster_dimension_probe_reports_dimensions() {
        // A valid PNG goes through the guarded decoder and reports its size.
        let img = image::DynamicImage::ImageRgba8(image::RgbaImage::from_fn(4, 3, |x, y| {
            image::Rgba([x as u8 * 60, y as u8 * 80, 0, 255])
        }));
        let mut png = Cursor::new(Vec::new());
        img.write_to(&mut png, ImageFormat::Png).unwrap();
        assert_eq!(
            inspect_image_dimensions("image/png", &png.into_inner()).unwrap(),
            (4, 3)
        );
    }

    #[test]
    fn raster_dimension_probe_is_guard_limited() {
        // A raster whose declared width is one pixel over the source cap must be
        // rejected by the probe's decompression-bomb guard, not decoded. Encodes
        // a real image so the rejection is attributable to `image::Limits` (the
        // width cap) rather than to malformed bytes.
        let img = image::DynamicImage::ImageRgba8(image::RgbaImage::new(
            choreo_image::MAX_SOURCE_DIMENSION + 1,
            1,
        ));
        let mut png = Cursor::new(Vec::new());
        img.write_to(&mut png, ImageFormat::Png).unwrap();
        assert!(inspect_image_dimensions("image/png", &png.into_inner()).is_err());
    }
}