import copy

import pytest

from litellm.proxy.common_utils.openai_endpoint_utils import (
    remove_sensitive_info_from_deployment,
)


@pytest.mark.parametrize(
    "model_config, expected_config",
    [
        # Test case 1: Empty litellm_params
        (
            {"model_name": "test-model", "litellm_params": {}},
            {"model_name": "test-model", "litellm_params": {}},
        ),
        # Test case 2: Full sensitive data removal, mixed secrets of azure, aws, gcp, and typical api_key
        (
            {
                "model_name": "gpt-4",
                "litellm_params": {
                    "model": "openai/gpt-4",
                    "api_key": "sk-sensitive-key-123",
                    "client_secret": "~v8Q4W:Zp9gJ-3sTqX5aB@LkR2mNfYdC",
                    "vertex_credentials": {"type": "service_account"},
                    "aws_access_key_id": "AKIA123456789",
                    "aws_secret_access_key": "secret-access-key",
                    "api_base": "https://api.openai.com/v1",
                    "temperature": 0.7,
                },
                "model_info": {"id": "test-id"},
            },
            {
                "model_name": "gpt-4",
                "litellm_params": {
                    "model": "openai/gpt-4",
                    "api_base": "https://api.openai.com/v1",
                    "temperature": 0.7,
                },
                "model_info": {"id": "test-id"},
            },
        ),
        # Test case 3: Partial sensitive data, api_key
        (
            {
                "model_name": "claude-3",
                "litellm_params": {
                    "model": "anthropic/claude-3",
                    "api_key": "sk-anthropic-key",
                    "temperature": 0.5,
                },
            },
            {
                "model_name": "claude-3",
                "litellm_params": {"model": "anthropic/claude-3", "temperature": 0.5},
            },
        ),
        # Test case 4: No sensitive data
        (
            {
                "model_name": "local-model",
                "litellm_params": {
                    "model": "local/model",
                    "temperature": 0.8,
                    "max_tokens": 100,
                },
            },
            {
                "model_name": "local-model",
                "litellm_params": {
                    "model": "local/model",
                    "temperature": 0.8,
                    "max_tokens": 100,
                },
            },
        ),
    ],
)
def test_remove_sensitive_info_from_deployment(
    model_config: dict, expected_config: dict
):
    sanitized_config = remove_sensitive_info_from_deployment(model_config)
    assert sanitized_config == expected_config


def test_remove_sensitive_info_from_deployment_with_excluded_keys():
    """
    Test that excluded_keys prevents masking of specific keys (exact match).
    """
    base_config = {
        "model_name": "test-model",
        "litellm_params": {
            "model": "openai/gpt-4",
            "api_key": "sk-sensitive-key-123",
            "litellm_credentials_name": "my-credential-name",
            "access_token": "token-12345",
            "temperature": 0.7,
        },
    }

    # Without excluded_keys, access_token should be masked (contains "token")
    sanitized_config = remove_sensitive_info_from_deployment(copy.deepcopy(base_config))
    assert sanitized_config["litellm_params"]["access_token"] != "token-12345"
    assert "*" in sanitized_config["litellm_params"]["access_token"]

    # With excluded_keys, litellm_credentials_name should NOT be masked.
    # ``remove_sensitive_info_from_deployment`` mutates its input, so feed it
    # a fresh copy rather than the already-sanitized one.
    sanitized_config = remove_sensitive_info_from_deployment(
        copy.deepcopy(base_config), excluded_keys={"litellm_credentials_name"}
    )
    assert (
        sanitized_config["litellm_params"]["litellm_credentials_name"]
        == "my-credential-name"
    )

    # access_token should still be masked (not in excluded_keys)
    assert sanitized_config["litellm_params"]["access_token"] != "token-12345"
    assert "*" in sanitized_config["litellm_params"]["access_token"]

    # api_key should still be removed (popped) regardless of excluded_keys
    assert "api_key" not in sanitized_config["litellm_params"]
