llm_providers 0.11.0

use serde::Serialize;

#[derive(Debug, Serialize, Clone, Copy)]
pub struct Model {
    pub id: &'static str,
    pub name: &'static str,
    pub description: &'static str,
    pub supports_tools: bool,
    pub context_length: Option<u64>,
    pub input_price: f64,
    pub output_price: f64,
    pub published_at: Option<&'static str>,
}

#[derive(Debug, Serialize, Clone, Copy)]
pub struct Endpoint {
    pub label: &'static str,
    pub region: &'static str,
    pub base_url: &'static str,
    pub price_currency: &'static str,
    pub docs_url: Option<&'static str>,
}

#[derive(Debug, Serialize)]
pub struct Provider {
    pub label: &'static str,
    pub endpoints: &'static phf::Map<&'static str, Endpoint>,
    pub models: &'static [Model],
    pub endpoint_models: Option<&'static phf::Map<&'static str, &'static [Model]>>,
}

include!(concat!(env!("OUT_DIR"), "/registry_generated.rs"));

pub type Providers = phf::Map<&'static str, Provider>;

/// Get the providers map
pub fn get_providers_data() -> &'static Providers {
    &PROVIDERS
}

pub fn registry_version() -> &'static str {
    REGISTRY_VERSION
}

pub fn registry_updated_at() -> &'static str {
    REGISTRY_UPDATED_AT
}

/// List all provider family IDs (sorted)
pub fn list_providers() -> Vec<String> {
    let mut keys: Vec<String> = get_providers_data()
        .keys()
        .cloned()
        .map(|s| s.to_string())
        .collect();
    keys.sort();
    keys
}

/// List all endpoint IDs (sorted)
pub fn list_endpoints() -> Vec<String> {
    let mut ids: Vec<String> = get_providers_data()
        .entries()
        .flat_map(|(provider_id, p)| {
            p.endpoints
                .keys()
                .cloned()
                .map(move |endpoint_key| format!("{}:{}", provider_id, endpoint_key))
        })
        .collect();
    ids.sort();
    ids
}

/// Find endpoint details by endpoint ID, returning (family_id, Endpoint)
///
/// Endpoint ID format: "{provider_id}:{endpoint_key}".
pub fn get_endpoint(endpoint_id: &str) -> Option<(&'static str, &'static Endpoint)> {
    if let Some((provider_id, endpoint_key)) = endpoint_id.split_once(':') {
        for (&pid, provider) in get_providers_data() {
            if pid == provider_id {
                let ep = provider.endpoints.get(endpoint_key)?;
                return Some((pid, ep));
            }
        }
        return None;
    }

    let mut found: Option<(&'static str, &'static Endpoint)> = None;
    for (&family_id, provider) in get_providers_data() {
        if let Some(ep) = provider.endpoints.get(endpoint_id) {
            if found.is_some() {
                return None;
            }
            found = Some((family_id, ep));
        }
    }
    found
}

/// List all model IDs under a provider family
pub fn list_models(provider_id: &str) -> Option<Vec<String>> {
    get_providers_data()
        .get(provider_id)
        .map(|p| p.models.iter().map(|m| m.id.to_string()).collect())
}

/// List all model IDs under a specific endpoint.
///
/// Endpoint ID format: "{provider_id}:{endpoint_key}".
///
/// Behavior:
/// - If the provider has endpoint-specific models, return those.
/// - Otherwise fall back to provider-level models.
pub fn list_models_for_endpoint(endpoint_id: &str) -> Option<Vec<String>> {
    let (provider_id, endpoint_key) = endpoint_id.split_once(':')?;
    let provider = get_providers_data().get(provider_id)?;

    if let Some(endpoint_models) = provider.endpoint_models {
        if let Some(models) = endpoint_models.get(endpoint_key) {
            return Some(models.iter().map(|m| m.id.to_string()).collect());
        }
    }

    Some(provider.models.iter().map(|m| m.id.to_string()).collect())
}

/// Get model details by (provider_id, model_id)
pub fn get_model(provider_id: &str, model_id: &str) -> Option<Model> {
    get_model_ref(provider_id, model_id).copied()
}

pub fn get_model_ref(provider_id: &str, model_id: &str) -> Option<&'static Model> {
    get_providers_data()
        .get(provider_id)
        .and_then(|p| p.models.iter().find(|m| m.id == model_id))
}

/// Get model details by endpoint_id and model_id.
///
/// Endpoint ID format: "{provider_id}:{endpoint_key}".
///
/// Behavior:
/// - If endpoint-specific models exist for that endpoint, search within them.
/// - Otherwise fall back to provider-level models.
pub fn get_model_for_endpoint(endpoint_id: &str, model_id: &str) -> Option<Model> {
    let (provider_id, endpoint_key) = endpoint_id.split_once(':')?;
    let provider = get_providers_data().get(provider_id)?;

    if let Some(endpoint_models) = provider.endpoint_models {
        if let Some(models) = endpoint_models.get(endpoint_key) {
            return models.iter().find(|m| m.id == model_id).copied();
        }
    }

    provider.models.iter().find(|m| m.id == model_id).copied()
}

#[derive(Default, Clone, Copy, PartialEq, Eq)]
pub enum OrderBy {
    #[default]
    ProviderAndModelId,
    PublishedAtAsc,
    PublishedAtDesc,
}

#[derive(Default)]
pub struct ModelFilter {
    pub provider_id: Option<String>,
    pub region: Option<String>,
    pub supports_tools: Option<bool>,
    pub min_context_length: Option<u64>,
    pub order_by: Option<OrderBy>,
}

/// Advanced filtering: returns a list of (provider_id, Model)
pub fn filter_models(filter: ModelFilter) -> Vec<(String, Model)> {
    let mut results = Vec::new();

    for (&pid, provider) in get_providers_data() {
        // Filter by provider_id
        if let Some(ref target_pid) = filter.provider_id {
            if pid != target_pid.as_str() {
                continue;
            }
        }

        // Filter by region: match if any endpoint has the target region
        if let Some(ref target_region) = filter.region {
            let has_region = provider
                .endpoints
                .values()
                .any(|ep| ep.region == target_region);
            if !has_region {
                continue;
            }
        }

        for model in provider.models {
            // Filter by supports_tools
            if let Some(target_tools) = filter.supports_tools {
                if model.supports_tools != target_tools {
                    continue;
                }
            }

            // Filter by min_context_length
            if let Some(min_ctx) = filter.min_context_length {
                if model.context_length.unwrap_or(0) < min_ctx {
                    continue;
                }
            }

            results.push((pid.to_string(), *model));
        }
    }

    // Sort results based on order_by
    match filter.order_by {
        Some(OrderBy::PublishedAtAsc) => {
            results.sort_by(|a, b| {
                match (a.1.published_at, b.1.published_at) {
                    (Some(a_date), Some(b_date)) => a_date.cmp(b_date),
                    (Some(_), None) => std::cmp::Ordering::Less,
                    (None, Some(_)) => std::cmp::Ordering::Greater,
                    (None, None) => {
                        let p_cmp = a.0.cmp(&b.0);
                        if p_cmp == std::cmp::Ordering::Equal {
                            a.1.id.cmp(b.1.id)
                        } else {
                            p_cmp
                        }
                    }
                }
            });
        }
        Some(OrderBy::PublishedAtDesc) => {
            results.sort_by(|a, b| {
                match (a.1.published_at, b.1.published_at) {
                    (Some(a_date), Some(b_date)) => b_date.cmp(a_date),
                    (Some(_), None) => std::cmp::Ordering::Less,
                    (None, Some(_)) => std::cmp::Ordering::Greater,
                    (None, None) => {
                        let p_cmp = a.0.cmp(&b.0);
                        if p_cmp == std::cmp::Ordering::Equal {
                            a.1.id.cmp(b.1.id)
                        } else {
                            p_cmp
                        }
                    }
                }
            });
        }
        _ => {
            // Default: sort by provider_id, then model_id
            results.sort_by(|a, b| {
                let p_cmp = a.0.cmp(&b.0);
                if p_cmp == std::cmp::Ordering::Equal {
                    a.1.id.cmp(b.1.id)
                } else {
                    p_cmp
                }
            });
        }
    }

    results
}

pub fn filter_models_ref(filter: ModelFilter) -> Vec<(&'static str, &'static Model)> {
    let mut results = Vec::new();

    for (&pid, provider) in get_providers_data() {
        if let Some(ref target_pid) = filter.provider_id {
            if pid != target_pid.as_str() {
                continue;
            }
        }

        if let Some(ref target_region) = filter.region {
            let has_region = provider
                .endpoints
                .values()
                .any(|ep| ep.region == target_region);
            if !has_region {
                continue;
            }
        }

        for model in provider.models {
            if let Some(target_tools) = filter.supports_tools {
                if model.supports_tools != target_tools {
                    continue;
                }
            }

            if let Some(min_ctx) = filter.min_context_length {
                if model.context_length.unwrap_or(0) < min_ctx {
                    continue;
                }
            }

            results.push((pid, model));
        }
    }

    // Sort results based on order_by
    match filter.order_by {
        Some(OrderBy::PublishedAtAsc) => {
            results.sort_by(|a, b| {
                match (a.1.published_at, b.1.published_at) {
                    (Some(a_date), Some(b_date)) => a_date.cmp(b_date),
                    (Some(_), None) => std::cmp::Ordering::Less,
                    (None, Some(_)) => std::cmp::Ordering::Greater,
                    (None, None) => {
                        let p_cmp = a.0.cmp(&b.0);
                        if p_cmp == std::cmp::Ordering::Equal {
                            a.1.id.cmp(b.1.id)
                        } else {
                            p_cmp
                        }
                    }
                }
            });
        }
        Some(OrderBy::PublishedAtDesc) => {
            results.sort_by(|a, b| {
                match (a.1.published_at, b.1.published_at) {
                    (Some(a_date), Some(b_date)) => b_date.cmp(a_date),
                    (Some(_), None) => std::cmp::Ordering::Less,
                    (None, Some(_)) => std::cmp::Ordering::Greater,
                    (None, None) => {
                        let p_cmp = a.0.cmp(&b.0);
                        if p_cmp == std::cmp::Ordering::Equal {
                            a.1.id.cmp(b.1.id)
                        } else {
                            p_cmp
                        }
                    }
                }
            });
        }
        _ => {
            // Default: sort by provider_id, then model_id
            results.sort_by(|a, b| {
                let p_cmp = a.0.cmp(&b.0);
                if p_cmp == std::cmp::Ordering::Equal {
                    a.1.id.cmp(b.1.id)
                } else {
                    p_cmp
                }
            });
        }
    }

    results
}

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

    #[test]
    fn test_parse_providers() {
        let providers = get_providers_data();
        assert!(!providers.is_empty());
        assert!(providers.contains_key("openai"));

        assert_eq!(registry_version(), "2.0");
    }

    #[test]
    fn test_provider_endpoints() {
        let providers = get_providers_data();
        let aliyun = providers.get("aliyun").expect("aliyun not found");
        let ep = aliyun
            .endpoints
            .get("cn")
            .expect("aliyun endpoint not found");
        assert_eq!(ep.region, "cn");
        assert_eq!(ep.price_currency, "CNY");

        // Multi-endpoint provider
        let moonshot = providers.get("moonshot").expect("moonshot not found");
        assert!(moonshot.endpoints.contains_key("cn"));
        assert!(moonshot.endpoints.contains_key("global"));
        assert_eq!(moonshot.endpoints["cn"].region, "cn");
        assert_eq!(moonshot.endpoints["global"].region, "global");
    }

    #[test]
    fn test_list_providers() {
        let providers = list_providers();
        assert!(providers.contains(&"openai".to_string()));
        assert!(providers.contains(&"anthropic".to_string()));
        // moonshot_global should NOT be a top-level provider anymore
        assert!(!providers.contains(&"moonshot_global".to_string()));
    }

    #[test]
    fn test_list_endpoints() {
        let endpoints = list_endpoints();
        assert!(endpoints.contains(&"openai:global".to_string()));
        assert!(endpoints.contains(&"moonshot:cn".to_string()));
        assert!(endpoints.contains(&"moonshot:global".to_string()));
    }

    #[test]
    fn test_get_endpoint() {
        let (family_id, ep) = get_endpoint("moonshot:global").expect("endpoint not found");
        assert_eq!(family_id, "moonshot");
        assert_eq!(ep.region, "global");
        assert_eq!(ep.price_currency, "USD");
    }

    #[test]
    fn test_list_models() {
        let models = list_models("openai").expect("OpenAI provider not found");
        assert!(models.contains(&"gpt-4o".to_string()));
    }

    #[test]
    fn test_get_model() {
        let model = get_model("openai", "gpt-4o").expect("Model not found");
        assert_eq!(model.id, "gpt-4o");
        assert!(model.supports_tools);
    }

    #[test]
    fn test_filter_models() {
        // 1. Filter by region="cn" — matches providers that have any CN endpoint
        let cn_models = filter_models(ModelFilter {
            region: Some("cn".to_string()),
            ..Default::default()
        });
        assert!(!cn_models.is_empty());
        assert!(cn_models.iter().any(|(p, _)| p == "aliyun"));
        assert!(!cn_models.iter().any(|(p, _)| p == "openai"));

        // 2. Filter by supports_tools=true
        let tool_models = filter_models(ModelFilter {
            supports_tools: Some(true),
            provider_id: Some("openai".to_string()),
            ..Default::default()
        });
        assert!(tool_models.iter().any(|(_, m)| m.id == "gpt-4o"));
    }

    #[tokio::test]
    async fn test_integration_with_llm_connector() {
        let providers = get_providers_data();
        if let Some(openai) = providers.get("openai") {
            let ep = openai
                .endpoints
                .get("global")
                .expect("openai endpoint not found");
            assert!(ep.base_url.contains("api.openai.com"));
            let has_gpt4o = openai.models.iter().any(|m| m.id == "gpt-4o");
            assert!(has_gpt4o, "OpenAI provider should have gpt-4o");
        }
    }

    #[test]
    fn test_filter_models_order_by_published_at_desc() {
        // Test sorting by published_at descending (newest first)
        let models = filter_models(ModelFilter {
            order_by: Some(OrderBy::PublishedAtDesc),
            ..Default::default()
        });

        // Find indices of models with known published_at dates
        let gpt55_idx = models.iter().position(|(_, m)| m.id == "gpt-5.5");
        let _gpt52_idx = models.iter().position(|(_, m)| m.id == "gpt-5.2");
        let _gemini3_idx = models.iter().position(|(_, m)| m.id == "gemini-3");
        let gemini25pro_idx = models.iter().position(|(_, m)| m.id == "gemini-2.5-pro");

        // Models with published_at should come before models without
        // And among models with dates, newer ones should come first
        if let (Some(idx1), Some(idx2)) = (gpt55_idx, gemini25pro_idx) {
            // gpt-5.5 (2025-03-01) is newer than gemini-2.5-pro (2024-09-24)
            assert!(idx1 < idx2, "gpt-5.5 should come before gemini-2.5-pro when sorting desc");
        }
    }

    #[test]
    fn test_filter_models_order_by_published_at_asc() {
        // Test sorting by published_at ascending (oldest first)
        let models = filter_models(ModelFilter {
            order_by: Some(OrderBy::PublishedAtAsc),
            ..Default::default()
        });

        // Find indices of models with known published_at dates
        let gpt55_idx = models.iter().position(|(_, m)| m.id == "gpt-5.5");
        let gemini25pro_idx = models.iter().position(|(_, m)| m.id == "gemini-2.5-pro");

        // Among models with dates, older ones should come first
        if let (Some(idx1), Some(idx2)) = (gemini25pro_idx, gpt55_idx) {
            // gemini-2.5-pro (2024-09-24) is older than gpt-5.5 (2025-03-01)
            assert!(idx1 < idx2, "gemini-2.5-pro should come before gpt-5.5 when sorting asc");
        }
    }

    #[test]
    fn test_model_published_at_field() {
        // Test that the published_at field is correctly loaded from JSON
        let gpt55 = get_model("openai", "gpt-5.5");
        assert!(gpt55.is_some(), "gpt-5.5 should exist");
        let model = gpt55.unwrap();
        assert_eq!(model.published_at, Some("2025-03-01"));
    }
}