imagegen-bridge-core 0.1.0

//! Common image-generation parameters and their wire representation.

use std::{fmt, str::FromStr};

use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de};

use crate::{BridgeError, ErrorCode};

macro_rules! string_enum {
    ($(#[$meta:meta])* $visibility:vis enum $name:ident {
        $($(#[$variant_meta:meta])* $variant:ident => $wire:literal),+ $(,)?
    }) => {
        $(#[$meta])*
        #[derive(
            Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash,
            Serialize, Deserialize, JsonSchema,
        )]
        #[serde(rename_all = "snake_case")]
        $visibility enum $name {
            $(
                $(#[$variant_meta])*
                #[doc = concat!("Wire value `", $wire, "`.")]
                #[serde(rename = $wire)]
                $variant
            ),+
        }

        impl fmt::Display for $name {
            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                formatter.write_str(match self { $(Self::$variant => $wire),+ })
            }
        }
    };
}

string_enum! {
    /// Requested generation quality.
    pub enum Quality {
        Auto => "auto",
        Low => "low",
        Medium => "medium",
        High => "high",
    }
}

impl Default for Quality {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Encoded output image format.
    pub enum OutputFormat {
        Png => "png",
        Jpeg => "jpeg",
        Webp => "webp",
    }
}

impl Default for OutputFormat {
    fn default() -> Self {
        Self::Png
    }
}

string_enum! {
    /// Requested background behavior.
    pub enum Background {
        Auto => "auto",
        Opaque => "opaque",
        Transparent => "transparent",
    }
}

string_enum! {
    /// How a transparent-background result is produced.
    pub enum TransparencyMode {
        /// Prefer native alpha and otherwise use a local chroma-key matte.
        Auto => "auto",
        /// Require provider-native alpha output.
        Native => "native",
        /// Generate a flat key background and remove it locally.
        ChromaKey => "chroma_key",
    }
}

impl Default for TransparencyMode {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Conditions under which an ordered provider fallback may run.
    pub enum FallbackPolicy {
        /// Fall back only for unavailable, unsupported, or pre-output failures.
        OnUnavailable => "on_unavailable",
        /// Also fall back for known-outcome provider failures, excluding safety and cancellation.
        OnError => "on_error",
    }
}

impl Default for FallbackPolicy {
    fn default() -> Self {
        Self::OnUnavailable
    }
}

string_enum! {
    /// Preferred execution shape for bridge-emulated multi-image fan-out.
    pub enum BatchExecution {
        /// Use provider/session-aware bridge defaults.
        Auto => "auto",
        /// Dispatch one fan-out chunk at a time.
        Sequential => "sequential",
        /// Use the provider's configured bounded parallelism.
        Parallel => "parallel",
    }
}

impl Default for BatchExecution {
    fn default() -> Self {
        Self::Auto
    }
}

impl Default for Background {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Provider moderation strictness where configurable.
    pub enum Moderation {
        Auto => "auto",
        Low => "low",
    }
}

impl Default for Moderation {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Behavior when one output in a multi-image request fails.
    pub enum MultiImageFailurePolicy {
        FailFast => "fail_fast",
        BestEffort => "best_effort",
    }
}

impl Default for MultiImageFailurePolicy {
    fn default() -> Self {
        Self::FailFast
    }
}

string_enum! {
    /// Fidelity used when processing image inputs.
    pub enum InputFidelity {
        Low => "low",
        High => "high",
    }
}

string_enum! {
    /// Image-generation tool action for conversational transports.
    pub enum ImageAction {
        Auto => "auto",
        Generate => "generate",
        Edit => "edit",
    }
}

impl Default for ImageAction {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Desired output payload representation.
    pub enum ResponseFormat {
        B64Json => "b64_json",
        Url => "url",
        Artifact => "artifact",
        Metadata => "metadata",
    }
}

impl Default for ResponseFormat {
    fn default() -> Self {
        Self::B64Json
    }
}

string_enum! {
    /// Atomic behavior when an explicit artifact filename already exists.
    pub enum ArtifactCollisionPolicy {
        Error => "error",
        Suffix => "suffix",
    }
}

impl Default for ArtifactCollisionPolicy {
    fn default() -> Self {
        Self::Error
    }
}

string_enum! {
    /// Optional portable generation-metadata persistence policy.
    pub enum ArtifactMetadataPolicy {
        None => "none",
        Sidecar => "sidecar",
        Embedded => "embedded",
        SidecarAndEmbedded => "sidecar_and_embedded",
    }
}

impl Default for ArtifactMetadataPolicy {
    fn default() -> Self {
        Self::None
    }
}

impl ArtifactMetadataPolicy {
    /// Whether the policy writes a JSON file beside each bridge-owned artifact.
    #[must_use]
    pub const fn writes_sidecar(self) -> bool {
        matches!(self, Self::Sidecar | Self::SidecarAndEmbedded)
    }

    /// Whether the policy embeds a bounded XMP generation record in image bytes.
    #[must_use]
    pub const fn embeds(self) -> bool {
        matches!(self, Self::Embedded | Self::SidecarAndEmbedded)
    }
}

string_enum! {
    /// Compatibility policy used during provider negotiation.
    pub enum CompatibilityMode {
        Strict => "strict",
        Normalize => "normalize",
        BestEffort => "best_effort",
    }
}

impl Default for CompatibilityMode {
    fn default() -> Self {
        Self::Strict
    }
}

string_enum! {
    /// Handling policy for a negative prompt.
    pub enum NegativePromptMode {
        Auto => "auto",
        Native => "native",
        Merge => "merge",
        Reject => "reject",
    }
}

impl Default for NegativePromptMode {
    fn default() -> Self {
        Self::Auto
    }
}

string_enum! {
    /// Visibility and requirement policy for an upstream revised prompt.
    pub enum RevisedPromptPolicy {
        Include => "include",
        Omit => "omit",
        Require => "require",
    }
}

impl Default for RevisedPromptPolicy {
    fn default() -> Self {
        Self::Include
    }
}

string_enum! {
    /// Codex conversation behavior for a request.
    pub enum SessionMode {
        Isolated => "isolated",
        Persistent => "persistent",
        Thread => "thread",
    }
}

impl Default for SessionMode {
    fn default() -> Self {
        Self::Isolated
    }
}

string_enum! {
    /// Coarse output resolution hint.
    pub enum Resolution {
        OneK => "1k",
        TwoK => "2k",
        FourK => "4k",
    }
}

/// Image size represented as `auto` or `WIDTHxHEIGHT`.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ImageSize(String);

impl ImageSize {
    /// Automatic provider-selected size.
    pub const AUTO: &'static str = "auto";

    /// Constructs a validated explicit size.
    pub fn exact(width: u32, height: u32) -> Result<Self, BridgeError> {
        if width == 0 || height == 0 {
            return Err(BridgeError::new(
                ErrorCode::InvalidRequest,
                "image dimensions must be greater than zero",
            ));
        }
        Ok(Self(format!("{width}x{height}")))
    }

    /// Returns `None` for `auto`, otherwise the explicit dimensions.
    #[must_use]
    pub fn dimensions(&self) -> Option<(u32, u32)> {
        if self.0 == Self::AUTO {
            return None;
        }
        let (width, height) = self.0.split_once('x')?;
        Some((width.parse().ok()?, height.parse().ok()?))
    }

    /// Returns true when the provider should choose the size.
    #[must_use]
    pub fn is_auto(&self) -> bool {
        self.0 == Self::AUTO
    }

    /// Returns the stable wire value.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl Default for ImageSize {
    fn default() -> Self {
        Self(Self::AUTO.to_owned())
    }
}

impl fmt::Display for ImageSize {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl FromStr for ImageSize {
    type Err = BridgeError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        if value == Self::AUTO {
            return Ok(Self::default());
        }
        let (width, height) = value.split_once('x').ok_or_else(|| {
            BridgeError::new(
                ErrorCode::InvalidRequest,
                "size must be `auto` or `WIDTHxHEIGHT`",
            )
        })?;
        if width.is_empty()
            || height.is_empty()
            || (width.len() > 1 && width.starts_with('0'))
            || (height.len() > 1 && height.starts_with('0'))
            || !width.bytes().all(|byte| byte.is_ascii_digit())
            || !height.bytes().all(|byte| byte.is_ascii_digit())
        {
            return Err(BridgeError::new(
                ErrorCode::InvalidRequest,
                "size must be `auto` or `WIDTHxHEIGHT`",
            ));
        }
        Self::exact(
            width.parse().map_err(|_| {
                BridgeError::new(ErrorCode::InvalidRequest, "image width is out of range")
            })?,
            height.parse().map_err(|_| {
                BridgeError::new(ErrorCode::InvalidRequest, "image height is out of range")
            })?,
        )
    }
}

impl Serialize for ImageSize {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for ImageSize {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = String::deserialize(deserializer)?;
        value.parse().map_err(de::Error::custom)
    }
}

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

    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
        json_schema!({
            "type": "string",
            "pattern": "^(auto|[1-9][0-9]*x[1-9][0-9]*)$",
            "examples": ["auto", "1024x1024", "1536x1024"]
        })
    }
}

/// Aspect ratio represented as `WIDTH:HEIGHT` with non-zero integer terms.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct AspectRatio(String);

impl AspectRatio {
    /// Constructs a reduced aspect ratio.
    pub fn new(width: u32, height: u32) -> Result<Self, BridgeError> {
        if width == 0 || height == 0 {
            return Err(BridgeError::new(
                ErrorCode::InvalidRequest,
                "aspect ratio terms must be greater than zero",
            ));
        }
        let divisor = gcd(width, height);
        Ok(Self(format!("{}:{}", width / divisor, height / divisor)))
    }

    /// Returns the reduced integer ratio.
    #[must_use]
    pub fn terms(&self) -> (u32, u32) {
        let (width, height) = self.0.split_once(':').unwrap_or(("1", "1"));
        (width.parse().unwrap_or(1), height.parse().unwrap_or(1))
    }

    /// Returns the stable wire value.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl FromStr for AspectRatio {
    type Err = BridgeError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        let (width, height) = value.split_once(':').ok_or_else(|| {
            BridgeError::new(
                ErrorCode::InvalidRequest,
                "aspect_ratio must use `WIDTH:HEIGHT`",
            )
        })?;
        Self::new(
            width.parse().map_err(|_| {
                BridgeError::new(ErrorCode::InvalidRequest, "aspect ratio width is invalid")
            })?,
            height.parse().map_err(|_| {
                BridgeError::new(ErrorCode::InvalidRequest, "aspect ratio height is invalid")
            })?,
        )
    }
}

impl fmt::Display for AspectRatio {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl Serialize for AspectRatio {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for AspectRatio {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?
            .parse()
            .map_err(de::Error::custom)
    }
}

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

    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
        json_schema!({
            "type": "string",
            "pattern": "^[1-9][0-9]*:[1-9][0-9]*$",
            "examples": ["1:1", "3:2", "16:9"]
        })
    }
}

const fn gcd(mut left: u32, mut right: u32) -> u32 {
    while right != 0 {
        let remainder = left % right;
        left = right;
        right = remainder;
    }
    left
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used)]

    use super::*;

    #[test]
    fn image_size_round_trips_as_a_string() {
        let size: ImageSize = "1536x1024".parse().unwrap();
        assert_eq!(size.dimensions(), Some((1536, 1024)));
        assert_eq!(serde_json::to_string(&size).unwrap(), "\"1536x1024\"");
        assert_eq!(
            serde_json::from_str::<ImageSize>("\"1536x1024\"").unwrap(),
            size
        );
    }

    #[test]
    fn image_size_rejects_ambiguous_values() {
        for value in ["", "1024", "0x1024", "1024X1024", "1x2x3", "01x1"] {
            assert!(value.parse::<ImageSize>().is_err(), "accepted {value}");
        }
    }

    #[test]
    fn aspect_ratio_is_reduced() {
        let ratio: AspectRatio = "1920:1080".parse().unwrap();
        assert_eq!(ratio.as_str(), "16:9");
        assert_eq!(ratio.terms(), (16, 9));
    }
}