use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_6_1_SOL_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_6_LUNA_MODEL_ID;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_6_SOL_MODEL_ID;
use codex_models_manager::bundled_models_response;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelVisibility;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::openai_models::WebSearchToolType;

const GPT_5_BEDROCK_CONTEXT_WINDOW: i64 = 272_000;
const GPT_5_6_SOL_OPENAI_MODEL_ID: &str = "gpt-5.6-sol";
const GPT_5_6_TERRA_OPENAI_MODEL_ID: &str = "gpt-5.6-terra";
const GPT_5_6_LUNA_OPENAI_MODEL_ID: &str = "gpt-5.6-luna";
const GPT_6_1_SOL_OPENAI_MODEL_ID: &str = "gpt-6.1-sol";
const GPT_6_SOL_OPENAI_MODEL_ID: &str = "gpt-6-sol";
const GPT_6_LUNA_OPENAI_MODEL_ID: &str = "gpt-6-luna";
const GPT_6_ASTRA_OPENAI_MODEL_ID: &str = "gpt-6-astra";
const GPT_5_5_OPENAI_MODEL_ID: &str = "gpt-5.5";

pub(crate) fn static_model_catalog() -> ModelsResponse {
    normalize_bedrock_catalog(ModelsResponse {
        models: vec![
            bedrock_model(
                bundled_openai_model(GPT_6_1_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_1_SOL_MODEL_ID,
                "GPT-6.1 Sol",
                /*priority*/ 0,
            ),
            bedrock_model(
                bundled_openai_model(GPT_6_ASTRA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID,
                "GPT-6-Astra",
                /*priority*/ 1,
            ),
            bedrock_model(
                bundled_openai_model(GPT_6_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_SOL_MODEL_ID,
                "GPT-6 Sol",
                /*priority*/ 2,
            ),
            bedrock_model(
                bundled_openai_model(GPT_6_LUNA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_LUNA_MODEL_ID,
                "GPT-6 Luna",
                /*priority*/ 3,
            ),
            bedrock_model(
                bundled_openai_model(GPT_5_6_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID,
                "GPT-5.6 Sol",
                /*priority*/ 4,
            ),
            bedrock_model(
                bundled_openai_model(GPT_5_6_TERRA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID,
                "GPT-5.6 Terra",
                /*priority*/ 5,
            ),
            bedrock_model(
                bundled_openai_model(GPT_5_6_LUNA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID,
                "GPT-5.6 Luna",
                /*priority*/ 6,
            ),
            gpt_5_bedrock_model(
                GPT_5_5_OPENAI_MODEL_ID,
                AMAZON_BEDROCK_GPT_5_5_MODEL_ID,
                "GPT-5.5",
                /*priority*/ 7,
            ),
        ],
    })
}

pub(super) fn static_gov_model_catalog() -> ModelsResponse {
    let mut catalog = static_model_catalog();
    catalog.models.retain(|model| {
        matches!(
            model.slug.as_str(),
            AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID | AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID
        )
    });
    catalog
}

pub(crate) fn normalize_bedrock_catalog(mut catalog: ModelsResponse) -> ModelsResponse {
    for model in &mut catalog.models {
        // Amazon Bedrock currently only supports the implicit "default" tier for GPT models.
        model.additional_speed_tiers.clear();
        model.service_tiers.clear();
        model.default_service_tier = None;
        // Bedrock rejects the `search_content_types` field used by multimodal search.
        model.web_search_tool_type = WebSearchToolType::Text;
    }
    catalog
}

fn gpt_5_bedrock_model(
    openai_slug: &str,
    bedrock_slug: &str,
    display_name: &str,
    priority: i32,
) -> ModelInfo {
    let mut model = bundled_openai_model(openai_slug);
    model.slug = bedrock_slug.to_string();
    model.display_name = display_name.to_string();
    model.priority = priority;
    model.context_window = Some(GPT_5_BEDROCK_CONTEXT_WINDOW);
    model.max_context_window = Some(GPT_5_BEDROCK_CONTEXT_WINDOW);
    model.visibility = ModelVisibility::List;
    model.availability_nux = None;
    model.upgrade = None;
    model
}

fn bedrock_model(
    mut model: ModelInfo,
    bedrock_slug: &str,
    display_name: &str,
    priority: i32,
) -> ModelInfo {
    model.slug = bedrock_slug.to_string();
    model.display_name = display_name.to_string();
    model.priority = priority;
    model.visibility = ModelVisibility::List;
    model.availability_nux = None;
    model.upgrade = None;
    model.use_responses_lite = false;
    model.tool_mode = None;
    model
}

fn bundled_openai_model(slug: &str) -> ModelInfo {
    bundled_models_response()
        .unwrap_or_else(|err| panic!("bundled models.json should parse: {err}"))
        .models
        .into_iter()
        .find(|model| model.slug == slug)
        .unwrap_or_else(|| panic!("bundled models.json should include {slug}"))
}

#[cfg(test)]
mod tests {
    use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE;
    use codex_protocol::protocol::MultiAgentVersion;
    use pretty_assertions::assert_eq;

    use super::*;

    #[test]
    fn catalog_uses_mantle_model_ids_in_priority_order() {
        let catalog = static_model_catalog();

        assert_eq!(
            catalog
                .models
                .iter()
                .map(|model| model.slug.as_str())
                .collect::<Vec<_>>(),
            vec![
                AMAZON_BEDROCK_GPT_6_1_SOL_MODEL_ID,
                AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID,
                AMAZON_BEDROCK_GPT_6_SOL_MODEL_ID,
                AMAZON_BEDROCK_GPT_6_LUNA_MODEL_ID,
                AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID,
                AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID,
                AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID,
                AMAZON_BEDROCK_GPT_5_5_MODEL_ID,
            ]
        );
    }

    #[test]
    fn gpt_5_bedrock_models_use_bedrock_context_window() {
        let catalog = static_model_catalog();

        assert_eq!(
            catalog
                .models
                .iter()
                .map(|model| (
                    model.slug.as_str(),
                    model.context_window,
                    model.max_context_window,
                    model.web_search_tool_type,
                ))
                .collect::<Vec<_>>(),
            vec![
                (
                    AMAZON_BEDROCK_GPT_6_1_SOL_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_6_SOL_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_6_LUNA_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(872_000),
                    WebSearchToolType::Text,
                ),
                (
                    AMAZON_BEDROCK_GPT_5_5_MODEL_ID,
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    Some(GPT_5_BEDROCK_CONTEXT_WINDOW),
                    WebSearchToolType::Text,
                ),
            ]
        );
    }

    #[test]
    fn configured_bedrock_catalogs_normalize_unsupported_model_capabilities() {
        let models = [
            Some(MultiAgentVersion::V2),
            Some(MultiAgentVersion::V1),
            Some(MultiAgentVersion::Disabled),
            None,
        ]
        .into_iter()
        .map(|version| {
            let mut model = bundled_openai_model(GPT_5_5_OPENAI_MODEL_ID);
            model.multi_agent_version = version;
            model
        })
        .collect::<Vec<_>>();
        let mut expected = models.clone();
        for model in &mut expected {
            model.additional_speed_tiers.clear();
            model.service_tiers.clear();
            model.default_service_tier = None;
            model.web_search_tool_type = WebSearchToolType::Text;
        }

        assert_eq!(
            normalize_bedrock_catalog(ModelsResponse { models }),
            ModelsResponse { models: expected }
        );
    }

    #[test]
    fn gpt_5_bedrock_models_do_not_include_availability_nux_or_upgrade() {
        let catalog = static_model_catalog();

        for model in catalog.models {
            assert_eq!((model.availability_nux, model.upgrade), (None, None));
        }
    }

    #[test]
    fn gpt_5_bedrock_models_are_visible() {
        let catalog = static_model_catalog();

        for model in catalog.models {
            assert_eq!(model.visibility, ModelVisibility::List);
        }
    }

    #[test]
    fn bedrock_models_preserve_source_metadata_with_supported_capabilities() {
        let catalog = static_model_catalog();

        for (mut expected, slug, display_name, priority) in [
            (
                bundled_openai_model(GPT_6_1_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_1_SOL_MODEL_ID,
                "GPT-6.1 Sol",
                0,
            ),
            (
                bundled_openai_model(GPT_6_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_SOL_MODEL_ID,
                "GPT-6 Sol",
                2,
            ),
            (
                bundled_openai_model(GPT_6_LUNA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_LUNA_MODEL_ID,
                "GPT-6 Luna",
                3,
            ),
            (
                bundled_openai_model(GPT_5_6_SOL_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID,
                "GPT-5.6 Sol",
                4,
            ),
            (
                bundled_openai_model(GPT_5_6_TERRA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID,
                "GPT-5.6 Terra",
                5,
            ),
            (
                bundled_openai_model(GPT_5_6_LUNA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID,
                "GPT-5.6 Luna",
                6,
            ),
            (
                bundled_openai_model(GPT_6_ASTRA_OPENAI_MODEL_ID),
                AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID,
                "GPT-6-Astra",
                1,
            ),
        ] {
            expected.slug = slug.to_string();
            expected.display_name = display_name.to_string();
            expected.priority = priority;
            expected.visibility = ModelVisibility::List;
            expected.availability_nux = None;
            expected.upgrade = None;
            expected.use_responses_lite = false;
            expected.tool_mode = None;
            expected.additional_speed_tiers.clear();
            expected.service_tiers.clear();
            expected.default_service_tier = None;
            expected.web_search_tool_type = WebSearchToolType::Text;

            assert_eq!(
                catalog.models.iter().find(|model| model.slug == slug),
                Some(&expected)
            );
        }
    }

    #[test]
    fn gpt_5_bedrock_models_only_allow_default_service_tier() {
        let catalog = static_model_catalog();

        for model in catalog.models {
            assert_eq!(model.additional_speed_tiers, Vec::<String>::new());
            assert_eq!(model.service_tiers, Vec::new());
            assert_eq!(model.default_service_tier, None);
            assert_eq!(
                model.service_tier_for_request(Some("priority".to_string())),
                None
            );
            assert_eq!(
                model
                    .service_tier_for_request(Some(SERVICE_TIER_DEFAULT_REQUEST_VALUE.to_string())),
                None
            );
        }
    }
}