litellm-rs 0.9.0

//! Gemini Model Registry
//!
//! Unified model registry system containing capabilities and pricing information for all Gemini models

use std::collections::HashMap;
use std::sync::OnceLock;

mod catalog;
#[cfg(any(test, feature = "providers-extended", feature = "providers-extra"))]
mod contract;
#[cfg(test)]
mod cost_tests;
#[cfg(test)]
mod current_tests;
mod pricing_schedule;
mod surface;

#[cfg(any(feature = "providers-extended", feature = "providers-extra"))]
pub(crate) use contract::{has_trailing_assistant_prefill, uses_fixed_sampling_contract};
#[cfg(any(feature = "providers-extended", feature = "providers-extra"))]
pub(crate) use pricing_schedule::flash_uses_standard_pricing_at;
pub use surface::GoogleGeminiApiSurface;

pub use crate::core::cost::types::ModelPricing;
use crate::core::types::model::ModelInfo;

/// Model features
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ModelFeature {
    /// Multimodal support (images, videos, audio)
    MultimodalSupport,
    /// Tool calling support
    ToolCalling,
    /// Function calling support
    FunctionCalling,
    /// Streaming support
    StreamingSupport,
    /// Context caching support
    ContextCaching,
    /// System instructions support
    SystemInstructions,
    /// Batch processing support
    BatchProcessing,
    /// JSON mode support
    JsonMode,
    /// Code execution support
    CodeExecution,
    /// Search grounding support
    SearchGrounding,
    /// Video understanding support
    VideoUnderstanding,
    /// Audio understanding support  
    AudioUnderstanding,
    /// Real-time streaming support
    RealtimeStreaming,
}

#[derive(Debug, Clone, PartialEq)]
pub enum GeminiModelFamily {
    Gemini37Flash,
    Gemini38Flash,
    Gemini36Flash,
    /// Gemini 3.5 series (2026 - Latest)
    Gemini35Flash,
    Gemini35FlashLite,

    Gemini31ProPreview,
    Gemini31Flash,
    Gemini31FlashLite,

    Gemini3Pro,
    Gemini3ProDeepThink,
    Gemini3Flash,
    Gemini3ProImage,

    Gemini25Pro,
    Gemini25Flash,
    Gemini25FlashLite,

    Gemini20Flash,
    Gemini20FlashThinking,

    Gemini15Pro,
    Gemini15Flash,
    Gemini15Flash8B,

    Gemini10Pro,
    Gemini10ProVision,

    GeminiExperimental,
}

fn pricing_per_million(
    input_price: f64,
    output_price: f64,
    cached_input_price: Option<f64>,
    image_price: Option<f64>,
    video_price_per_second: Option<f64>,
    audio_price_per_second: Option<f64>,
) -> ModelPricing {
    ModelPricing {
        input_cost_per_1k_tokens: input_price / 1000.0,
        output_cost_per_1k_tokens: output_price / 1000.0,
        cache_read_input_token_cost: cached_input_price.map(|price| price / 1000.0),
        cost_per_image: image_price.map(|price| {
            let mut costs = HashMap::new();
            costs.insert("default".to_string(), price);
            costs
        }),
        video_cost_per_second: video_price_per_second,
        audio_cost_per_second: audio_price_per_second,
        currency: "USD".to_string(),
        updated_at: chrono::Utc::now(),
        ..Default::default()
    }
}

/// Model limits
#[derive(Debug, Clone)]
pub struct ModelLimits {
    /// Maximum context length
    pub max_context_length: u32,
    /// Maximum output tokens
    pub max_output_tokens: u32,
    /// Maximum image count
    pub max_images: Option<u32>,
    /// Maximum video length (seconds)
    pub max_video_seconds: Option<u32>,
    /// Maximum audio length (seconds)
    pub max_audio_seconds: Option<u32>,
    /// Requests per minute limit
    pub rpm_limit: Option<u32>,
    /// Tokens per minute limit
    pub tpm_limit: Option<u32>,
}

/// Model
#[derive(Debug, Clone)]
pub struct ModelSpec {
    /// Model
    pub model_info: ModelInfo,
    /// Model
    pub family: GeminiModelFamily,
    /// Supported features
    pub features: Vec<ModelFeature>,
    /// Pricing information
    pub pricing: ModelPricing,
    /// Limit information
    pub limits: ModelLimits,
}

/// Model
#[derive(Debug, Clone)]
pub struct GeminiModelRegistry {
    models: HashMap<String, ModelSpec>,
}

impl GeminiModelRegistry {
    /// Expected number of Gemini models for capacity hint
    const EXPECTED_MODEL_COUNT: usize = 20;

    pub fn new() -> Self {
        let mut registry = Self {
            models: HashMap::with_capacity(Self::EXPECTED_MODEL_COUNT),
        };
        registry.initialize_models();
        registry
    }

    fn initialize_models(&mut self) {
        catalog::register_all(self);
    }

    fn register_model(&mut self, id: &str, mut spec: ModelSpec) {
        spec.pricing.model = id.to_string();
        self.models.insert(id.to_string(), spec);
    }

    pub fn get_model_spec(&self, model_id: &str) -> Option<&ModelSpec> {
        self.models.get(model_id)
    }

    pub fn list_models(&self) -> Vec<&ModelSpec> {
        self.models.values().collect()
    }

    pub fn supports_feature(&self, model_id: &str, feature: &ModelFeature) -> bool {
        self.get_model_spec(model_id)
            .map(|spec| spec.features.contains(feature))
            .unwrap_or(false)
    }

    pub fn get_model_family(&self, model_id: &str) -> Option<&GeminiModelFamily> {
        self.get_model_spec(model_id).map(|spec| &spec.family)
    }

    /// Model
    pub fn get_model_pricing(&self, model_id: &str) -> Option<&ModelPricing> {
        self.get_model_spec(model_id).map(|spec| &spec.pricing)
    }

    /// Get model pricing in the shared core cost model shape.
    pub fn get_core_model_pricing(&self, model_id: &str) -> Option<ModelPricing> {
        self.get_model_spec(model_id)
            .map(|spec| spec.pricing.clone())
    }

    /// Model
    pub fn get_model_limits(&self, model_id: &str) -> Option<&ModelLimits> {
        self.get_model_spec(model_id).map(|spec| &spec.limits)
    }

    /// Detect model family from model name string
    pub fn from_model_name(model_name: &str) -> Option<GeminiModelFamily> {
        let model_lower = model_name.to_lowercase();

        if model_name == "gemini-3.8-flash" {
            Some(GeminiModelFamily::Gemini38Flash)
        } else if model_name == "gemini-3.7-flash" {
            Some(GeminiModelFamily::Gemini37Flash)
        } else if model_lower.contains("gemini-3.6-flash") {
            Some(GeminiModelFamily::Gemini36Flash)
        } else if model_lower.contains("gemini-3.5-flash-lite") {
            Some(GeminiModelFamily::Gemini35FlashLite)
        } else if model_lower.contains("gemini-3.5-flash") {
            Some(GeminiModelFamily::Gemini35Flash)
        }
        // Gemini 3.1 series (check before 3.0 as more specific)
        else if model_lower.contains("gemini-3.1-flash-lite") {
            Some(GeminiModelFamily::Gemini31FlashLite)
        } else if model_lower.contains("gemini-3.1-flash") {
            Some(GeminiModelFamily::Gemini31Flash)
        } else if model_lower.contains("gemini-3.1-pro") {
            Some(GeminiModelFamily::Gemini31ProPreview)
        }
        // Gemini 3.0 series (deprecated 2026-03-09)
        else if model_lower.contains("gemini-3") && model_lower.contains("deep-think") {
            Some(GeminiModelFamily::Gemini3ProDeepThink)
        } else if model_lower.contains("gemini-3") && model_lower.contains("image") {
            Some(GeminiModelFamily::Gemini3ProImage)
        } else if model_lower.contains("gemini-3-flash") || model_lower.contains("gemini-3.0-flash")
        {
            Some(GeminiModelFamily::Gemini3Flash)
        } else if model_lower.contains("gemini-3-pro") || model_lower.contains("gemini-3.0-pro") {
            Some(GeminiModelFamily::Gemini3Pro)
        }
        // Gemini 2.5 series
        else if model_lower.contains("gemini-2.5-flash-lite") {
            Some(GeminiModelFamily::Gemini25FlashLite)
        } else if model_lower.contains("gemini-2.5-flash") {
            Some(GeminiModelFamily::Gemini25Flash)
        } else if model_lower.contains("gemini-2.5-pro") {
            Some(GeminiModelFamily::Gemini25Pro)
        }
        // Gemini 2.0 series
        else if model_lower.contains("gemini-2.0-flash-thinking") {
            Some(GeminiModelFamily::Gemini20FlashThinking)
        } else if model_lower.contains("gemini-2.0-flash") || model_lower.contains("gemini-2-flash")
        {
            Some(GeminiModelFamily::Gemini20Flash)
        }
        // Gemini 1.5 series
        else if model_lower.contains("gemini-1.5-pro") || model_lower.contains("gemini-15-pro") {
            Some(GeminiModelFamily::Gemini15Pro)
        } else if model_lower.contains("gemini-1.5-flash-8b") {
            Some(GeminiModelFamily::Gemini15Flash8B)
        } else if model_lower.contains("gemini-1.5-flash")
            || model_lower.contains("gemini-15-flash")
        {
            Some(GeminiModelFamily::Gemini15Flash)
        }
        // Gemini 1.0 series
        else if model_lower.contains("gemini-1.0-pro-vision") {
            Some(GeminiModelFamily::Gemini10ProVision)
        } else if model_lower.contains("gemini-1.0-pro") || model_lower.contains("gemini-pro") {
            Some(GeminiModelFamily::Gemini10Pro)
        }
        // Experimental
        else if model_lower.contains("gemini-exp") {
            Some(GeminiModelFamily::GeminiExperimental)
        } else {
            None
        }
    }
}

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

/// Model
pub fn get_gemini_registry() -> &'static GeminiModelRegistry {
    static REGISTRY: OnceLock<GeminiModelRegistry> = OnceLock::new();
    REGISTRY.get_or_init(GeminiModelRegistry::new)
}

/// Cost calculation utility
pub struct CostCalculator;

impl CostCalculator {
    /// Calculate basic cost
    pub fn calculate_cost(
        model_id: &str,
        prompt_tokens: u32,
        completion_tokens: u32,
    ) -> Result<f64, crate::ProviderError> {
        super::calculate_gemini_cost(model_id, prompt_tokens, completion_tokens)
    }

    /// Estimate multimodal cost. Prompt tokens include the audio input tokens.
    /// Audio duration uses Google's documented 32 tokens/second. Mixed cached
    /// audio cannot be priced from these arguments and returns `None`.
    pub fn calculate_multimodal_cost(
        model_id: &str,
        prompt_tokens: u32,
        completion_tokens: u32,
        cached_tokens: Option<u32>,
        images: Option<u32>,
        video_seconds: Option<u32>,
        audio_seconds: Option<u32>,
    ) -> Option<f64> {
        let registry = get_gemini_registry();
        let pricing = registry.get_core_model_pricing(model_id)?;
        let mut usage =
            crate::core::cost::types::UsageTokens::new(prompt_tokens, completion_tokens);
        usage.cached_tokens = cached_tokens;
        let mut total_cost =
            match crate::core::cost::calculator::generic_cost_per_token(model_id, &usage, "gemini")
            {
                Ok(cost) => cost.total_cost,
                Err(crate::core::cost::types::CostError::ModelNotSupported { .. }) => {
                    let rates = [
                        pricing.input_cost_per_1k_tokens,
                        pricing.output_cost_per_1k_tokens,
                        pricing.cache_read_input_token_cost.unwrap_or(0.0),
                    ];
                    if rates.iter().any(|rate| !rate.is_finite() || *rate < 0.0)
                        || rates.iter().all(|rate| *rate == 0.0)
                    {
                        return None;
                    }
                    let cached_tokens = cached_tokens.unwrap_or(0);
                    let remaining_prompt_tokens = prompt_tokens.saturating_sub(cached_tokens);
                    let cache_cost = pricing
                        .cache_read_input_token_cost
                        .map(|price| cached_tokens as f64 / 1000.0 * price)
                        .unwrap_or(0.0);
                    cache_cost
                        + remaining_prompt_tokens as f64 / 1000.0 * pricing.input_cost_per_1k_tokens
                        + completion_tokens as f64 / 1000.0 * pricing.output_cost_per_1k_tokens
                }
                Err(_) => return None,
            };

        // Image cost
        let image_price = pricing
            .cost_per_image
            .as_ref()
            .and_then(|costs| costs.get("default"))
            .copied();
        if let (Some(img_count), Some(img_price)) = (images, image_price) {
            total_cost += img_count as f64 * img_price;
        }

        // Video cost
        if let (Some(video_secs), Some(video_price)) =
            (video_seconds, pricing.video_cost_per_second)
        {
            total_cost += video_secs as f64 * video_price;
        }

        // Audio is included in prompt_tokens: replace its text-rate portion,
        // rather than adding the full audio rate on top of already billed input.
        // https://ai.google.dev/gemini-api/docs/audio#technical-details-about-audio
        if let Some(seconds) = audio_seconds.filter(|seconds| *seconds > 0) {
            if let Some(audio_rate) = pricing.input_cost_per_audio_token {
                let audio_tokens = seconds.checked_mul(32)?;
                if audio_tokens > prompt_tokens || cached_tokens.unwrap_or(0) > 0 {
                    return None;
                }
                total_cost +=
                    audio_tokens as f64 * (audio_rate - pricing.input_cost_per_1k_tokens / 1000.0);
            } else if let Some(audio_rate) = pricing.audio_cost_per_second {
                total_cost += seconds as f64 * audio_rate;
            }
        }

        Some(total_cost)
    }

    /// Estimate token count
    pub fn estimate_tokens(text: &str) -> u32 {
        // Gemini uses approximately 4 characters = 1 token ratio (English)
        (text.len() as f32 / 4.0).ceil() as u32
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::types::model::ProviderCapability;

    #[test]
    fn test_model_registry() {
        let registry = get_gemini_registry();

        // Test Gemini 2.5 Flash
        let flash_spec = registry.get_model_spec("gemini-2.5-flash").unwrap();
        assert_eq!(flash_spec.family, GeminiModelFamily::Gemini25Flash);
        assert!(
            flash_spec
                .features
                .contains(&ModelFeature::MultimodalSupport)
        );
        assert!(
            flash_spec
                .features
                .contains(&ModelFeature::VideoUnderstanding)
        );

        // Test pricing
        assert_eq!(flash_spec.pricing.input_cost_per_1k_tokens, 0.0003);
        assert_eq!(flash_spec.pricing.output_cost_per_1k_tokens, 0.0025);
    }

    #[test]
    fn test_model_family_detection() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-3.5-flash"),
            Some(GeminiModelFamily::Gemini35Flash)
        );

        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-2.0-flash-exp"),
            Some(GeminiModelFamily::Gemini20Flash)
        );

        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-3.1-flash"),
            Some(GeminiModelFamily::Gemini31Flash)
        );

        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-1.5-pro-latest"),
            Some(GeminiModelFamily::Gemini15Pro)
        );

        assert_eq!(GeminiModelRegistry::from_model_name("unknown-model"), None);
    }

    #[test]
    fn test_cost_calculation() {
        let cost = CostCalculator::calculate_cost("gemini-2.5-flash", 1000, 500);
        assert!(cost.is_ok());

        let cost_value = cost.unwrap();
        // Expected: (1000/1M * $0.30) + (500/1M * $2.50) = $0.0003 + $0.00125 = $0.00155
        assert!((cost_value - 0.00155).abs() < 0.000001);
    }

    #[test]
    fn test_cost_calculation_does_not_use_registry_fallback() {
        assert!(matches!(
            CostCalculator::calculate_cost("gemini-1.0-pro", 1000, 500),
            Err(crate::ProviderError::ModelNotFound { .. })
        ));
    }

    #[test]
    fn test_feature_support() {
        let registry = get_gemini_registry();

        // Gemini 2.5 Flash supports video understanding
        assert!(registry.supports_feature("gemini-2.5-flash", &ModelFeature::VideoUnderstanding));

        // Gemini 1.0 Pro does not support multimodal
        assert!(!registry.supports_feature("gemini-1.0-pro", &ModelFeature::VideoUnderstanding));
    }

    #[test]
    fn test_registry_default() {
        let registry = GeminiModelRegistry::default();
        assert!(!registry.models.is_empty());
    }

    #[test]
    fn test_list_models() {
        let registry = get_gemini_registry();
        let models = registry.list_models();
        assert!(!models.is_empty());
        assert!(models.len() >= 10); // We have at least 10 models registered
    }

    #[test]
    fn test_get_model_family() {
        let registry = get_gemini_registry();
        let family = registry.get_model_family("gemini-2.5-pro");
        assert!(family.is_some());
        assert_eq!(*family.unwrap(), GeminiModelFamily::Gemini25Pro);

        let family_unknown = registry.get_model_family("unknown-model");
        assert!(family_unknown.is_none());
    }

    #[test]
    fn test_get_model_pricing() {
        let registry = get_gemini_registry();
        let pricing = registry.get_model_pricing("gemini-2.5-flash");
        assert!(pricing.is_some());
        let pricing_value = pricing.unwrap();
        assert_eq!(pricing_value.input_cost_per_1k_tokens, 0.0003);
        assert_eq!(pricing_value.output_cost_per_1k_tokens, 0.0025);
        assert!(pricing_value.cache_read_input_token_cost.is_some());
    }

    #[test]
    fn test_core_model_pricing_conversion() {
        let registry = get_gemini_registry();
        let pricing = registry
            .get_core_model_pricing("gemini-2.5-flash")
            .expect("registry pricing should convert to core pricing");

        assert_eq!(pricing.model, "gemini-2.5-flash");
        assert_eq!(pricing.input_cost_per_1k_tokens, 0.0003);
        assert_eq!(pricing.output_cost_per_1k_tokens, 0.0025);
        assert!((pricing.cache_read_input_token_cost.unwrap() - 0.00003).abs() < 1e-12);
        assert_eq!(pricing.video_cost_per_second, None);
        assert_eq!(pricing.audio_cost_per_second, None);
        assert_eq!(
            pricing
                .cost_per_image
                .as_ref()
                .and_then(|costs| costs.get("default"))
                .copied(),
            None
        );
        assert_eq!(pricing.currency, "USD");
    }

    #[test]
    fn test_gemini_35_flash_metadata() {
        let registry = get_gemini_registry();
        let Some(spec) = registry.get_model_spec("gemini-3.5-flash") else {
            panic!("gemini-3.5-flash should be registered");
        };

        assert_eq!(spec.family, GeminiModelFamily::Gemini35Flash);
        assert_eq!(spec.limits.max_context_length, 1_048_576);
        assert_eq!(spec.limits.max_output_tokens, 65_536);
        assert_eq!(spec.limits.max_images, Some(3000));
        assert_eq!(spec.limits.max_video_seconds, Some(3600));
        assert_eq!(spec.limits.max_audio_seconds, Some(9600));
        assert_eq!(spec.limits.rpm_limit, None);
        assert_eq!(spec.limits.tpm_limit, None);
        assert!(spec.features.contains(&ModelFeature::ToolCalling));
        assert!(spec.features.contains(&ModelFeature::ContextCaching));
        assert!(spec.features.contains(&ModelFeature::SearchGrounding));
        assert_eq!(spec.pricing.input_cost_per_1k_tokens, 0.0015);
        assert_eq!(spec.pricing.output_cost_per_1k_tokens, 0.009);
        assert_eq!(spec.pricing.cache_read_input_token_cost, Some(0.00015));
        assert_eq!(spec.pricing.cost_per_image, None);
        assert_eq!(spec.pricing.video_cost_per_second, None);
        assert_eq!(spec.pricing.audio_cost_per_second, None);
    }

    #[test]
    fn test_gemini_family_capabilities_include_declared_advanced_features() {
        let registry = get_gemini_registry();
        let full_capability_models = [
            "gemini-2.5-pro",
            "gemini-2.5-flash",
            "gemini-2.5-flash-lite",
            "gemini-3-flash-preview",
            "gemini-3.1-pro-preview",
            "gemini-3.1-flash-lite",
        ];

        for model_id in full_capability_models {
            let Some(spec) = registry.get_model_spec(model_id) else {
                panic!("{model_id} should be registered");
            };
            for capability in [
                ProviderCapability::FunctionCalling,
                ProviderCapability::CodeExecution,
                ProviderCapability::BatchProcessing,
            ] {
                assert!(
                    spec.model_info.capabilities.contains(&capability),
                    "{model_id} should expose {capability:?}"
                );
            }
        }
    }

    #[test]
    fn test_get_model_limits() {
        let registry = get_gemini_registry();
        let limits = registry.get_model_limits("gemini-2.5-pro");
        assert!(limits.is_some());
        let limits_value = limits.unwrap();
        assert_eq!(limits_value.max_context_length, 1_048_576);
        assert_eq!(limits_value.max_output_tokens, 65536);
    }

    #[test]
    fn test_model_family_detection_gemini_3() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-3-pro"),
            Some(GeminiModelFamily::Gemini3Pro)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-3-pro-deep-think"),
            Some(GeminiModelFamily::Gemini3ProDeepThink)
        );
    }

    #[test]
    fn test_model_family_detection_gemini_25() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-2.5-pro"),
            Some(GeminiModelFamily::Gemini25Pro)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-2.5-flash"),
            Some(GeminiModelFamily::Gemini25Flash)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-2.5-flash-lite"),
            Some(GeminiModelFamily::Gemini25FlashLite)
        );
    }

    #[test]
    fn test_model_family_detection_gemini_20() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-2.0-flash-thinking-exp"),
            Some(GeminiModelFamily::Gemini20FlashThinking)
        );
    }

    #[test]
    fn test_model_family_detection_gemini_15() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-1.5-flash-8b"),
            Some(GeminiModelFamily::Gemini15Flash8B)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-1.5-flash"),
            Some(GeminiModelFamily::Gemini15Flash)
        );
    }

    #[test]
    fn test_model_family_detection_gemini_10() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-1.0-pro"),
            Some(GeminiModelFamily::Gemini10Pro)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-1.0-pro-vision"),
            Some(GeminiModelFamily::Gemini10ProVision)
        );
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-pro"),
            Some(GeminiModelFamily::Gemini10Pro)
        );
    }

    #[test]
    fn test_model_family_detection_experimental() {
        assert_eq!(
            GeminiModelRegistry::from_model_name("gemini-exp-something"),
            Some(GeminiModelFamily::GeminiExperimental)
        );
    }

    #[test]
    fn test_cost_calculation_unknown_model() {
        let cost = CostCalculator::calculate_cost("unknown-model", 1000, 500);
        assert!(matches!(
            cost,
            Err(crate::ProviderError::ModelNotFound { .. })
        ));
    }

    #[test]
    fn test_multimodal_cost_calculation() {
        let cost = CostCalculator::calculate_multimodal_cost(
            "gemini-2.5-flash",
            1000,
            500,
            Some(200),
            Some(5),
            None,
            None,
        );
        assert!(cost.is_some());
        let cost_value = cost.unwrap();
        assert!(cost_value > 0.0);
    }

    #[test]
    fn test_multimodal_cost_with_video_and_audio() {
        let cost = CostCalculator::calculate_multimodal_cost(
            "gemini-2.5-flash",
            5000,
            500,
            None,
            Some(5),
            Some(60),
            Some(120),
        );
        assert!(cost.is_some());
        let cost_value = cost.unwrap();
        assert!(cost_value > 0.0);
    }

    #[test]
    fn test_estimate_tokens() {
        let tokens = CostCalculator::estimate_tokens("Hello, world!");
        assert!((3..=5).contains(&tokens));
    }

    #[test]
    fn test_feature_support_unknown_model() {
        let registry = get_gemini_registry();
        assert!(!registry.supports_feature("unknown-model", &ModelFeature::MultimodalSupport));
    }

    #[test]
    fn test_gemini_25_pro_features() {
        let registry = get_gemini_registry();
        let spec = registry.get_model_spec("gemini-2.5-pro").unwrap();

        assert!(spec.features.contains(&ModelFeature::ToolCalling));
        assert!(spec.features.contains(&ModelFeature::FunctionCalling));
        assert!(spec.features.contains(&ModelFeature::StreamingSupport));
        assert!(spec.features.contains(&ModelFeature::ContextCaching));
        assert!(spec.features.contains(&ModelFeature::SystemInstructions));
        assert!(spec.features.contains(&ModelFeature::BatchProcessing));
        assert!(spec.features.contains(&ModelFeature::JsonMode));
        assert!(spec.features.contains(&ModelFeature::CodeExecution));
        assert!(spec.features.contains(&ModelFeature::SearchGrounding));
        assert!(spec.features.contains(&ModelFeature::VideoUnderstanding));
        assert!(spec.features.contains(&ModelFeature::AudioUnderstanding));
    }

    #[test]
    fn retired_models_are_not_callable() {
        let registry = get_gemini_registry();
        for model in [
            "gemini-1.0-pro",
            "gemini-1.5-pro",
            "gemini-1.5-flash",
            "gemini-1.5-flash-8b",
            "gemini-2.0-flash-exp",
            "gemini-2.0-flash-thinking-exp",
            "gemini-3-pro-image-preview",
            "gemini-3-pro",
            "gemini-3-pro-deep-think",
            "gemini-3.1-flash",
        ] {
            assert!(registry.get_model_spec(model).is_none(), "{model}");
            assert!(
                !registry.supports_feature(model, &ModelFeature::StreamingSupport),
                "{model}"
            );
        }
    }

    #[test]
    fn test_model_info_structure() {
        let registry = get_gemini_registry();
        let spec = registry.get_model_spec("gemini-2.5-flash").unwrap();

        assert_eq!(spec.model_info.id, "gemini-2.5-flash");
        assert_eq!(spec.model_info.provider, "gemini");
        assert!(spec.model_info.supports_streaming);
        assert!(spec.model_info.supports_tools);
        assert!(spec.model_info.supports_multimodal);
    }
}