import os
from unittest.mock import patch

import pytest

from litellm.llms.datarobot.chat.transformation import DataRobotConfig


@patch.dict(os.environ, {}, clear=True)
class TestDataRobotConfig:
    @pytest.fixture
    def handler(self):
        return DataRobotConfig()

    @pytest.mark.parametrize(
        "api_base, expected_url",
        [
            (None, "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"),
            (
                "http://localhost:5001",
                "http://localhost:5001/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com",
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com/api/v2/",
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com/api/v2",
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions",
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
                "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://staging.datarobot.com",
                "https://staging.datarobot.com/api/v2/genai/llmgw/chat/completions/",
            ),
            (
                "https://app.datarobot.com/api/v2/deployments/deployment_id",
                "https://app.datarobot.com/api/v2/deployments/deployment_id/",
            ),
            (
                "https://app.datarobot.com/api/v2/deployments/deployment_id/",
                "https://app.datarobot.com/api/v2/deployments/deployment_id/",
            ),
        ],
    )
    def test_resolve_api_base(self, api_base, expected_url, handler):
        """Test that URLs properly resolve to the expected format."""
        assert handler._resolve_api_base(api_base) == expected_url

        # Check that the complete url with the resolution is expected
        assert (
            handler.get_complete_url(
                api_base=handler._resolve_api_base(api_base),
                api_key="PASSTHROUGH_KEY",
                model="datarobot/vertex_ai/gemini-1.5-flash-002",
                optional_params={},
                litellm_params={},
            )
            == expected_url
        )

        # Check that the complete url with the original api_base does not change the url
        if api_base is not None:
            assert (
                handler.get_complete_url(
                    api_base=api_base,
                    api_key="PASSTHROUGH_KEY",
                    model="datarobot/vertex_ai/gemini-1.5-flash-002",
                    optional_params={},
                    litellm_params={},
                )
                == api_base
            )

    def test_resolve_api_base_with_environment_variable(self, handler, monkeypatch):
        monkeypatch.setenv("DATAROBOT_ENDPOINT", "https://env.datarobot.com")
        assert (
            handler._resolve_api_base(None)
            == "https://env.datarobot.com/api/v2/genai/llmgw/chat/completions/"
        )
        del os.environ["DATAROBOT_ENDPOINT"]

    @pytest.mark.parametrize(
        "api_key, expected_api_key",
        [
            (None, "fake-api-key"),
            ("PASSTHROUGH_KEY", "PASSTHROUGH_KEY"),
        ],
    )
    def test_resolve_api_key(self, api_key, expected_api_key, handler):
        assert handler._resolve_api_key(api_key) == expected_api_key

    def test_resolve_api_key_with_environment_variable(self, handler, monkeypatch):
        monkeypatch.setenv("DATAROBOT_API_TOKEN", "env_key")
        assert handler._resolve_api_key(None) == "env_key"
        del os.environ["DATAROBOT_API_TOKEN"]
