use super::types::{ModelInfo, ProviderInfo, ProvidersRegistry};

impl ProvidersRegistry {
    pub fn get_provider(&self, id: &str) -> Option<&ProviderInfo> {
        self.providers.get(id)
    }

    pub fn get_model(&self, provider: &str, model: &str) -> Option<&ModelInfo> {
        let mut result = self
            .providers
            .get(provider)
            .and_then(|provider| provider.models.get(model));

        // Fallback: if provider is "codex" and model not found, try "openai"
        if result.is_none() && provider == "codex" {
            result = self
                .providers
                .get("openai")
                .and_then(|provider| provider.models.get(model));
        }

        // Fallback: if provider is "kimi-code" and model not found, try "moonshotai".
        if result.is_none() && provider == "kimi-code" {
            result = self
                .providers
                .get("moonshotai")
                .and_then(|provider| provider.models.get(model));
        }

        // General fallback: progressively strip the last `-suffix` segment from
        // the provider name to find a parent entry.  This lets user-configured
        // aliases like "zai-coding" resolve model info from "zai" (or
        // "zai-coding-plan" from "zai-coding" -> "zai"), without needing
        // per-provider hardcoded rules.
        if result.is_none() {
            let mut candidate = provider;
            while let Some(pos) = candidate.rfind('-') {
                candidate = &candidate[..pos];
                if let Some(info) = self.providers.get(candidate)
                    && let Some(m) = info.models.get(model)
                {
                    result = Some(m);
                    break;
                }
            }
        }

        result
    }

    pub fn list_providers(&self) -> Vec<&str> {
        self.providers.keys().map(|s| s.as_str()).collect()
    }

    pub fn list_models(&self, provider: &str) -> Vec<&str> {
        self.providers
            .get(provider)
            .map(|provider| provider.models.keys().map(|s| s.as_str()).collect())
            .unwrap_or_default()
    }

    pub fn get_pricing(&self, provider: &str, model: &str) -> Option<&super::types::ModelPricing> {
        self.get_model(provider, model).map(|m| &m.pricing)
    }

    pub fn get_limits(&self, provider: &str, model: &str) -> Option<&super::types::ModelLimits> {
        self.get_model(provider, model).map(|m| &m.limits)
    }

    /// Backward-compatible alias for `get_limits`.
    pub fn get_constraints(
        &self,
        provider: &str,
        model: &str,
    ) -> Option<&super::types::ModelLimits> {
        self.get_limits(provider, model)
    }

    pub fn get_capabilities(
        &self,
        provider: &str,
        model: &str,
    ) -> Option<&super::types::ModelCapabilities> {
        self.get_model(provider, model).map(|m| &m.capabilities)
    }
}

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

    fn create_test_registry() -> ProvidersRegistry {
        let mut providers = HashMap::new();

        // Create openai provider with a test model
        let mut openai_models = HashMap::new();
        openai_models.insert(
            "gpt-4".to_string(),
            ModelInfo {
                id: "gpt-4".to_string(),
                name: "GPT-4".to_string(),
                ..Default::default()
            },
        );

        providers.insert(
            "openai".to_string(),
            ProviderInfo {
                id: "openai".to_string(),
                name: "OpenAI".to_string(),
                models: openai_models,
                ..Default::default()
            },
        );

        // Create codex provider without the model
        providers.insert(
            "codex".to_string(),
            ProviderInfo {
                id: "codex".to_string(),
                name: "Codex".to_string(),
                models: HashMap::new(),
                ..Default::default()
            },
        );

        // Create moonshotai provider with a test model
        let mut moonshot_models = HashMap::new();
        moonshot_models.insert(
            "kimi-k2".to_string(),
            ModelInfo {
                id: "kimi-k2".to_string(),
                name: "Kimi K2".to_string(),
                ..Default::default()
            },
        );

        providers.insert(
            "moonshotai".to_string(),
            ProviderInfo {
                id: "moonshotai".to_string(),
                name: "Moonshot".to_string(),
                models: moonshot_models,
                ..Default::default()
            },
        );

        // Create kimi-code provider without models to exercise fallback.
        providers.insert(
            "kimi-code".to_string(),
            ProviderInfo {
                id: "kimi-code".to_string(),
                name: "Kimi Code".to_string(),
                models: HashMap::new(),
                ..Default::default()
            },
        );

        ProvidersRegistry { providers }
    }

    #[test]
    fn test_get_model_direct_lookup() {
        let registry = create_test_registry();

        // Direct lookup should work
        assert!(registry.get_model("openai", "gpt-4").is_some());
    }

    #[test]
    fn test_get_model_codex_fallback() {
        let registry = create_test_registry();

        // Codex provider exists but model doesn't, should fallback to openai
        let model = registry.get_model("codex", "gpt-4");
        assert!(model.is_some());
        assert_eq!(model.unwrap().id, "gpt-4");
    }

    #[test]
    fn test_get_model_codex_no_fallback_when_found() {
        let mut registry = create_test_registry();

        // Add a model to codex
        let codex_provider = registry.providers.get_mut("codex").unwrap();
        codex_provider.models.insert(
            "codex-model".to_string(),
            ModelInfo {
                id: "codex-model".to_string(),
                name: "Codex Model".to_string(),
                ..Default::default()
            },
        );

        // Should find in codex, not fallback
        let model = registry.get_model("codex", "codex-model");
        assert!(model.is_some());
        assert_eq!(model.unwrap().id, "codex-model");
    }

    #[test]
    fn test_get_model_non_codex_no_fallback() {
        let registry = create_test_registry();

        // Non-codex provider shouldn't fallback
        assert!(registry.get_model("openai", "nonexistent").is_none());
        assert!(registry.get_model("other", "gpt-4").is_none());
    }

    #[test]
    fn test_get_model_kimi_code_fallback() {
        let registry = create_test_registry();

        let model = registry.get_model("kimi-code", "kimi-k2");
        assert!(model.is_some());
        assert_eq!(model.unwrap().id, "kimi-k2");
    }

    #[test]
    fn test_get_model_prefix_fallback_strips_suffix() {
        let mut registry = create_test_registry();

        // Add a "zai" provider with a model
        let mut zai_models = HashMap::new();
        zai_models.insert(
            "glm-5.1".to_string(),
            ModelInfo {
                id: "glm-5.1".to_string(),
                name: "GLM-5.1".to_string(),
                ..Default::default()
            },
        );
        registry.providers.insert(
            "zai".to_string(),
            ProviderInfo {
                id: "zai".to_string(),
                name: "Z.AI".to_string(),
                models: zai_models,
                ..Default::default()
            },
        );

        // "zai-coding" is not in registry, should fall back to "zai"
        let model = registry.get_model("zai-coding", "glm-5.1");
        assert!(model.is_some());
        assert_eq!(model.unwrap().id, "glm-5.1");
    }

    #[test]
    fn test_get_model_prefix_fallback_multi_level() {
        let mut registry = create_test_registry();

        let mut models = HashMap::new();
        models.insert(
            "model-a".to_string(),
            ModelInfo {
                id: "model-a".to_string(),
                name: "Model A".to_string(),
                ..Default::default()
            },
        );
        registry.providers.insert(
            "alibaba".to_string(),
            ProviderInfo {
                id: "alibaba".to_string(),
                name: "Alibaba".to_string(),
                models,
                ..Default::default()
            },
        );

        // "alibaba-coding-plan-cn" -> "alibaba-coding-plan" -> "alibaba-coding" -> "alibaba"
        let model = registry.get_model("alibaba-coding-plan-cn", "model-a");
        assert!(model.is_some());
        assert_eq!(model.unwrap().id, "model-a");
    }

    #[test]
    fn test_get_model_prefix_fallback_no_match() {
        let registry = create_test_registry();

        // "unknown-provider-xyz" should not match anything
        assert!(
            registry
                .get_model("unknown-provider-xyz", "gpt-4")
                .is_none()
        );
    }

    #[test]
    fn test_get_model_prefix_fallback_prefers_direct() {
        let mut registry = create_test_registry();

        // Add both "zai" and "zai-coding" with different models
        let mut zai_models = HashMap::new();
        zai_models.insert(
            "glm-5.1".to_string(),
            ModelInfo {
                id: "glm-5.1".to_string(),
                name: "GLM-5.1 General".to_string(),
                ..Default::default()
            },
        );
        registry.providers.insert(
            "zai".to_string(),
            ProviderInfo {
                id: "zai".to_string(),
                name: "Z.AI".to_string(),
                models: zai_models,
                ..Default::default()
            },
        );

        let mut coding_models = HashMap::new();
        coding_models.insert(
            "glm-5.1".to_string(),
            ModelInfo {
                id: "glm-5.1".to_string(),
                name: "GLM-5.1 Coding".to_string(),
                ..Default::default()
            },
        );
        registry.providers.insert(
            "zai-coding".to_string(),
            ProviderInfo {
                id: "zai-coding".to_string(),
                name: "Z.AI Coding".to_string(),
                models: coding_models,
                ..Default::default()
            },
        );

        // Direct lookup should win over fallback
        let model = registry.get_model("zai-coding", "glm-5.1");
        assert!(model.is_some());
        assert_eq!(model.unwrap().name, "GLM-5.1 Coding");
    }

    #[test]
    fn test_fallback_propagates_through_helper_methods() {
        let registry = create_test_registry();

        assert!(registry.get_pricing("codex", "gpt-4").is_some());
        assert!(registry.get_constraints("codex", "gpt-4").is_some());
        assert!(registry.get_capabilities("codex", "gpt-4").is_some());

        assert!(registry.get_pricing("kimi-code", "kimi-k2").is_some());
        assert!(registry.get_constraints("kimi-code", "kimi-k2").is_some());
        assert!(registry.get_capabilities("kimi-code", "kimi-k2").is_some());
    }
}
