litellm-rs 0.8.2

//! Fal AI Model Registry
//!
//! Model definitions and registry for Fal AI image generation models

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Image size representation for Fal AI
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ImageSize {
    /// Preset size name (e.g., "square_hd", "landscape_4_3")
    Preset(String),
    /// Custom dimensions
    Custom { width: u32, height: u32 },
}

impl ImageSize {
    /// Convert OpenAI size format (e.g., "1024x1024") to Fal AI format
    pub fn from_openai_size(size: &str) -> Self {
        // Map standard OpenAI sizes to Fal AI presets
        match size {
            "1024x1024" => ImageSize::Preset("square_hd".to_string()),
            "512x512" => ImageSize::Preset("square".to_string()),
            "1792x1024" => ImageSize::Preset("landscape_16_9".to_string()),
            "1024x1792" => ImageSize::Preset("portrait_16_9".to_string()),
            "1024x768" => ImageSize::Preset("landscape_4_3".to_string()),
            "768x1024" => ImageSize::Preset("portrait_4_3".to_string()),
            _ => {
                // Try to parse custom dimensions
                if let Some((w, h)) = size.split_once('x')
                    && let (Ok(width), Ok(height)) = (w.parse(), h.parse())
                {
                    return ImageSize::Custom { width, height };
                }
                // Default fallback
                ImageSize::Preset("landscape_4_3".to_string())
            }
        }
    }
}

/// Fal AI model definition
#[derive(Debug, Clone)]
pub struct FalAIModel {
    /// Model ID/endpoint (e.g., "fal-ai/flux/schnell")
    pub id: String,
    /// Display name
    pub name: String,
    /// Model description
    pub description: String,
    /// Fixed cost per image in USD; None when pricing needs additional request facts.
    pub cost_per_image: Option<f64>,
    /// Supported image sizes
    pub supported_sizes: Vec<String>,
    /// Maximum number of images per request
    pub max_images: u32,
    /// Whether the model supports prompt enhancement
    pub supports_prompt_enhancement: bool,
}

impl FalAIModel {
    /// Create a new Fal AI model definition
    pub fn new(id: &str, name: &str, description: &str, cost_per_image: f64) -> Self {
        Self {
            id: id.to_string(),
            name: name.to_string(),
            description: description.to_string(),
            cost_per_image: Some(cost_per_image),
            supported_sizes: vec![
                "square".to_string(),
                "square_hd".to_string(),
                "portrait_4_3".to_string(),
                "portrait_16_9".to_string(),
                "landscape_4_3".to_string(),
                "landscape_16_9".to_string(),
            ],
            max_images: 4,
            supports_prompt_enhancement: false,
        }
    }

    /// Set maximum images per request
    pub fn with_max_images(mut self, max: u32) -> Self {
        self.max_images = max;
        self
    }

    /// Set supported sizes
    pub fn with_sizes(mut self, sizes: Vec<String>) -> Self {
        self.supported_sizes = sizes;
        self
    }

    /// Enable prompt enhancement support
    pub fn with_prompt_enhancement(mut self) -> Self {
        self.supports_prompt_enhancement = true;
        self
    }
}

/// Fal AI model registry
#[derive(Debug, Clone)]
pub struct FalAIModelRegistry {
    models: HashMap<String, FalAIModel>,
}

impl Default for FalAIModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl FalAIModelRegistry {
    /// Create a new model registry with default models
    pub fn new() -> Self {
        let mut models = HashMap::new();

        // Official endpoint schemas and pricing cards, reviewed 2026-10-03.
        // None means the tariff needs facts unavailable to token-only pricing.
        let entries = [
            ("fal-ai/flux/schnell", "Flux Schnell", None, 4, false, false),
            (
                "fal-ai/flux-pro/v1.1",
                "Flux Pro v1.1",
                None,
                4,
                false,
                true,
            ),
            (
                "fal-ai/flux-pro/v1.1-ultra",
                "Flux Pro v1.1 Ultra",
                None,
                4,
                true,
                true,
            ),
            (
                "fal-ai/stable-diffusion-v3-medium",
                "Stable Diffusion 3 Medium",
                Some(0.035),
                4,
                false,
                false,
            ),
            (
                "fal-ai/recraft/v3/text-to-image",
                "Recraft V3",
                None,
                1,
                false,
                false,
            ),
            ("fal-ai/ideogram/v3", "Ideogram V3", None, 8, false, false),
            (
                "fal-ai/bria/text-to-image/hd",
                "BRIA HD",
                None,
                4,
                true,
                false,
            ),
            (
                "fal-ai/recraft/v4/text-to-image",
                "Recraft V4",
                Some(0.04),
                1,
                false,
                false,
            ),
            (
                "fal-ai/recraft/v4/pro/text-to-image",
                "Recraft V4 Pro",
                Some(0.25),
                1,
                false,
                false,
            ),
            ("fal-ai/flux-2-pro", "FLUX.2 Pro", None, 1, false, false),
            ("fal-ai/flux-2-flex", "FLUX.2 Flex", None, 1, false, false),
            ("ideogram/v4", "Ideogram V4", None, 4, false, false),
        ];
        for (id, name, fixed_cost, max_images, aspect_ratio_only, enhance_prompt) in entries {
            let mut model = FalAIModel::new(id, name, "Text-to-image generation", 0.0);
            model.cost_per_image = fixed_cost;
            model.max_images = max_images;
            model.supports_prompt_enhancement = enhance_prompt;
            if aspect_ratio_only {
                // These endpoints accept a ratio, not pixel sizes or presets.
                model.supported_sizes.clear();
            }
            models.insert(id.to_string(), model);
        }

        Self { models }
    }

    /// Get model by ID
    pub fn get(&self, model_id: &str) -> Option<&FalAIModel> {
        self.models.get(model_id)
    }

    /// Check if model is supported
    pub fn is_supported(&self, model_id: &str) -> bool {
        self.models.contains_key(model_id)
    }

    /// List all supported models
    pub fn list_models(&self) -> Vec<&FalAIModel> {
        self.models.values().collect()
    }

    /// Get cost per image for a model
    pub fn get_cost_per_image(&self, model_id: &str) -> Option<f64> {
        self.models.get(model_id).and_then(|m| m.cost_per_image)
    }

    /// Register a custom model
    pub fn register(&mut self, model: FalAIModel) {
        self.models.insert(model.id.clone(), model);
    }
}

/// Supported OpenAI parameters for image generation
pub const SUPPORTED_OPENAI_PARAMS: &[&str] = &["n", "response_format", "size"];

/// Map OpenAI parameters to Fal AI parameters
pub fn map_openai_to_fal_params(params: &serde_json::Value) -> serde_json::Value {
    let mut fal_params = serde_json::Map::new();

    if let Some(obj) = params.as_object() {
        for (key, value) in obj {
            match key.as_str() {
                "n" => {
                    fal_params.insert("num_images".to_string(), value.clone());
                }
                "response_format" => {
                    // Fal AI uses output_format
                    let format = match value.as_str() {
                        Some("b64_json") | Some("url") => "jpeg",
                        Some(f) => f,
                        None => "jpeg",
                    };
                    fal_params.insert("output_format".to_string(), serde_json::json!(format));
                }
                "size" => {
                    if let Some(size_str) = value.as_str() {
                        let image_size = ImageSize::from_openai_size(size_str);
                        fal_params.insert(
                            "image_size".to_string(),
                            serde_json::to_value(image_size).unwrap_or_default(),
                        );
                    }
                }
                _ => {
                    // Pass through other parameters
                    fal_params.insert(key.clone(), value.clone());
                }
            }
        }
    }

    serde_json::Value::Object(fal_params)
}

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

    #[test]
    fn test_image_size_from_openai_standard() {
        let size = ImageSize::from_openai_size("1024x1024");
        match size {
            ImageSize::Preset(s) => assert_eq!(s, "square_hd"),
            _ => panic!("Expected preset size"),
        }
    }

    #[test]
    fn test_image_size_from_openai_custom() {
        let size = ImageSize::from_openai_size("800x600");
        match size {
            ImageSize::Custom { width, height } => {
                assert_eq!(width, 800);
                assert_eq!(height, 600);
            }
            _ => panic!("Expected custom size"),
        }
    }

    #[test]
    fn test_image_size_from_openai_invalid() {
        let size = ImageSize::from_openai_size("invalid");
        match size {
            ImageSize::Preset(s) => assert_eq!(s, "landscape_4_3"),
            _ => panic!("Expected fallback preset"),
        }
    }

    #[test]
    fn test_model_registry_default() {
        let registry = FalAIModelRegistry::new();
        assert!(registry.is_supported("fal-ai/flux/schnell"));
        assert!(registry.is_supported("fal-ai/flux-pro/v1.1"));
    }

    #[test]
    fn test_model_registry_get() {
        let registry = FalAIModelRegistry::new();
        let model = registry.get("fal-ai/flux/schnell");
        assert!(model.is_some());
        assert_eq!(model.unwrap().name, "Flux Schnell");
    }

    #[test]
    fn test_model_registry_cost() {
        let registry = FalAIModelRegistry::new();
        let cost = registry.get_cost_per_image("fal-ai/stable-diffusion-v3-medium");
        assert!(cost.unwrap() > 0.0);
    }

    #[test]
    fn test_model_registry_unknown() {
        let registry = FalAIModelRegistry::new();
        assert!(!registry.is_supported("unknown-model"));
        assert_eq!(registry.get_cost_per_image("unknown-model"), None);
    }

    #[test]
    fn test_map_openai_to_fal_params() {
        let openai_params = serde_json::json!({
            "n": 2,
            "size": "1024x1024",
            "response_format": "url"
        });
        let fal_params = map_openai_to_fal_params(&openai_params);

        assert_eq!(fal_params["num_images"], 2);
        assert_eq!(fal_params["output_format"], "jpeg");
        assert!(fal_params.get("image_size").is_some());
    }

    #[test]
    fn test_model_with_builder() {
        let model = FalAIModel::new("test", "Test Model", "A test model", 0.01)
            .with_max_images(8)
            .with_prompt_enhancement();

        assert_eq!(model.max_images, 8);
        assert!(model.supports_prompt_enhancement);
    }

    #[test]
    fn test_register_custom_model() {
        let mut registry = FalAIModelRegistry::new();
        let custom = FalAIModel::new("custom/model", "Custom", "Custom model", 0.1);
        registry.register(custom);

        assert!(registry.is_supported("custom/model"));
    }

    #[test]
    fn test_list_models() {
        let registry = FalAIModelRegistry::new();
        let models = registry.list_models();
        assert!(!models.is_empty());
    }
}