Image Token Calculator — Rust source

//! Image Token Calculator — estimate the vision token cost of an image using
//! OpenAI-style tile math.
//!
//! Language: Rust (edition 2021, standard library only)
//! Source:   CosmoDev polyglot showcase port of the Image Token Calculator
//!           tool, ported from src/lib/imageTokenCalculator.ts (the canonical
//!           TypeScript implementation).
//! Live at:  https://dev.cosmolabs.org/tools/image-token-calculator
//! License:  display source — part of CosmoDev's polyglot tool pages.
//!
//! Design goals:
//!   - Pure + deterministic; invalid input is a `Result::Err`, never a panic.
//!   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//!   - Self-contained: std only (no crates.io dependencies).
//!
//! Rounding note: the TS reference uses Math.round (half up). Rust's `f64::
//! round` rounds half away from zero — identical on the positive values that
//! occur here — but `shrink` spells it as `floor(x + 0.5)` for exact parity.

/// Fixed token cost of the low-resolution image view.
pub const LOW_DETAIL_TOKENS: u32 = 85;
/// Token cost of one high-resolution 512 px tile.
pub const TILE_TOKENS: u32 = 170;
/// Images are first scaled to fit inside this square.
pub const MAX_SIDE: u32 = 2048;
/// Then the shortest side is capped at this length.
pub const MAX_SHORT_SIDE: u32 = 768;
/// Tile edge length in pixels.
pub const TILE_SIZE: u32 = 512;
/// Both dimensions at or under this → `Auto` stays low detail.
pub const AUTO_LOW_MAX: u32 = 512;

/// Requested detail mode of an image (`Auto` mirrors the TS default).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DetailLevel {
    Low,
    High,
    Auto,
}

/// Why an estimate was rejected: non-positive dimensions or an unknown
/// detail level (the TS reference throws for both).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ImageTokenError {
    /// Width/height must be integers greater than zero.
    InvalidDimensions,
    /// The detail level string was not "low", "high", or "auto".
    UnknownDetail(String),
}

impl std::fmt::Display for ImageTokenError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ImageTokenError::InvalidDimensions => {
                write!(f, "Width and height must be greater than zero")
            }
            ImageTokenError::UnknownDetail(d) => write!(f, "Unknown detail level: {d}"),
        }
    }
}

/// Mirrors the `TokenBreakdown` interface in the TS lib.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TokenBreakdown {
    /// Detail level actually applied (`Auto` resolves to low or high).
    pub detail: &'static str,
    /// Dimensions after the high-detail downscaling pipeline (identity for low).
    pub scaled_width: u32,
    pub scaled_height: u32,
    /// 512 px tiles along each axis (both 1 in low detail).
    pub tiles_x: u32,
    pub tiles_y: u32,
    /// Total 512 px tiles used (`tiles_x * tiles_y`).
    pub tiles: u32,
    /// Fixed base cost of the low-resolution view, in tokens.
    pub base: u32,
    /// Extra tokens for the high-resolution tile views (0 in low detail).
    pub detail_tokens: u32,
    /// Total estimated tokens: `base + detail_tokens`.
    pub total: u32,
}

/// JS `Math.round` parity, floored at 1 px: half up, never zero.
fn shrink(side: u32, scale: f64) -> u32 {
    ((side as f64 * scale) + 0.5).floor().max(1.0) as u32
}

/// Scale `(width, height)` per the vision preprocessing pipeline:
/// 1. fit inside a `MAX_SIDE` × `MAX_SIDE` square (longest side capped), then
/// 2. cap the shortest side at `MAX_SHORT_SIDE`.
///
/// Aspect ratio is preserved; each step is skipped when already satisfied.
pub fn preprocess_image(width: u32, height: u32) -> (u32, u32) {
    let mut w = width;
    let mut h = height;
    let longest = w.max(h);
    if longest > MAX_SIDE {
        let scale = MAX_SIDE as f64 / longest as f64;
        w = shrink(w, scale);
        h = shrink(h, scale);
    }
    let shortest = w.min(h);
    if shortest > MAX_SHORT_SIDE {
        let scale = MAX_SHORT_SIDE as f64 / shortest as f64;
        w = shrink(w, scale);
        h = shrink(h, scale);
    }
    (w, h)
}

/// Estimate the token cost of a `width` × `height` image at the given detail
/// level.
///
/// - [`DetailLevel::Low`]: fixed [`LOW_DETAIL_TOKENS`], whatever the size.
/// - [`DetailLevel::High`]: the image is downscaled by [`preprocess_image`],
///   tiled into [`TILE_SIZE`] squares, and each tile costs [`TILE_TOKENS`] on
///   top of the base.
/// - [`DetailLevel::Auto`]: low when both dimensions are ≤ [`AUTO_LOW_MAX`],
///   otherwise high.
pub fn image_tokens(
    width: u32,
    height: u32,
    detail: DetailLevel,
) -> Result<TokenBreakdown, ImageTokenError> {
    if width == 0 || height == 0 {
        return Err(ImageTokenError::InvalidDimensions);
    }

    let resolved_high = match detail {
        DetailLevel::Low => false,
        DetailLevel::High => true,
        DetailLevel::Auto => width > AUTO_LOW_MAX || height > AUTO_LOW_MAX,
    };

    if !resolved_high {
        return Ok(TokenBreakdown {
            detail: "low",
            scaled_width: width,
            scaled_height: height,
            tiles_x: 1,
            tiles_y: 1,
            tiles: 1,
            base: LOW_DETAIL_TOKENS,
            detail_tokens: 0,
            total: LOW_DETAIL_TOKENS,
        });
    }

    let (scaled_width, scaled_height) = preprocess_image(width, height);
    let tiles_x = scaled_width.div_ceil(TILE_SIZE);
    let tiles_y = scaled_height.div_ceil(TILE_SIZE);
    let tiles = tiles_x * tiles_y;
    let detail_tokens = tiles * TILE_TOKENS;
    Ok(TokenBreakdown {
        detail: "high",
        scaled_width,
        scaled_height,
        tiles_x,
        tiles_y,
        tiles,
        base: LOW_DETAIL_TOKENS,
        detail_tokens,
        total: LOW_DETAIL_TOKENS + detail_tokens,
    })
}

/// Parse a detail-level string ("low" | "high" | "auto") — the bridge from
/// the TS string union to [`DetailLevel`].
pub fn parse_detail(detail: &str) -> Result<DetailLevel, ImageTokenError> {
    match detail {
        "low" => Ok(DetailLevel::Low),
        "high" => Ok(DetailLevel::High),
        "auto" => Ok(DetailLevel::Auto),
        other => Err(ImageTokenError::UnknownDetail(other.to_string())),
    }
}