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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelRegistry {
    #[serde(flatten)]
    pub providers: HashMap<String, ProviderInfo>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderInfo {
    pub id: String,
    pub name: String,
    pub models: HashMap<String, ModelMetadata>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
    pub id: String,
    pub name: String,
    pub cost: Option<ModelCost>,
    pub limit: Option<ModelLimit>,
    pub modalities: Option<ModelModalities>,
    pub tool_call: Option<bool>,
    pub reasoning: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCost {
    pub input: f64,
    pub output: f64,
    pub cache_read: Option<f64>,
    pub cache_write: Option<f64>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelLimit {
    pub context: u32,
    pub output: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelModalities {
    pub input: Vec<String>,
    pub output: Vec<String>,
}

impl ModelRegistry {
    /// Find a model by its ID (searching across all providers)
    pub fn find_model(&self, model_id: &str) -> Option<&ModelMetadata> {
        // First try exact match if the key in models map matches the ID
        for provider in self.providers.values() {
            if let Some(model) = provider.models.get(model_id) {
                return Some(model);
            }
        }

        // Try searching for the model ID inside the metadata if the key was different
        for provider in self.providers.values() {
            for model in provider.models.values() {
                if model.id == model_id {
                    return Some(model);
                }
            }
        }

        None
    }
}
