import os
import sys
from unittest.mock import MagicMock, call, patch

import pytest

sys.path.insert(
    0, os.path.abspath("../../..")
)  # Adds the parent directory to the system path

from litellm.llms.watsonx.common_utils import generate_iam_token


class TestGenerateIAMToken:
    """Tests for the generate_iam_token function, specifically testing API key fallback logic."""

    @patch("litellm.llms.watsonx.common_utils.iam_token_cache")
    @patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
    @patch("litellm.llms.watsonx.common_utils.get_secret_str")
    def test_generate_iam_token_with_watsonx_zenapikey(
        self, mock_get_secret_str, mock_client, mock_cache
    ):
        """Test that WATSONX_ZENAPIKEY is used when it's the only key available."""
        # Setup mocks
        mock_cache.get_cache.return_value = None  # Cache miss
        mock_get_secret_str.side_effect = lambda key: (
            "zen-api-key-12345" if key == "WATSONX_ZENAPIKEY" else None
        )

        mock_response = MagicMock()
        mock_response.json.return_value = {
            "access_token": "test-token-12345",
            "expires_in": 3600,
        }
        mock_response.raise_for_status = MagicMock()
        mock_client.post.return_value = mock_response

        # Call function without api_key parameter
        result = generate_iam_token()

        # Verify get_secret_str was called with correct keys in order
        # Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL")
        calls = [
            call[0][0]
            for call in mock_get_secret_str.call_args_list
            if call[0][0] != "WATSONX_IAM_URL"
        ]
        assert "WX_API_KEY" in calls
        assert "WATSONX_API_KEY" in calls
        assert "WATSONX_APIKEY" in calls
        assert "WATSONX_ZENAPIKEY" in calls

        # Verify the token was generated using WATSONX_ZENAPIKEY
        assert result == "test-token-12345"
        mock_client.post.assert_called_once()
        call_kwargs = mock_client.post.call_args
        assert call_kwargs.kwargs["data"]["apikey"] == "zen-api-key-12345"

    @patch("litellm.llms.watsonx.common_utils.iam_token_cache")
    @patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
    @patch("litellm.llms.watsonx.common_utils.get_secret_str")
    def test_generate_iam_token_api_key_priority_order(
        self, mock_get_secret_str, mock_client, mock_cache
    ):
        """Test that API keys are checked in the correct priority order."""
        # Setup mocks
        mock_cache.get_cache.return_value = None  # Cache miss

        # Test priority: WX_API_KEY > WATSONX_API_KEY > WATSONX_APIKEY > WATSONX_ZENAPIKEY
        test_cases = [
            # (env_keys_set, expected_key_used, expected_calls)
            (
                {"WX_API_KEY": "wx-key"},
                "wx-key",
                ["WX_API_KEY"],  # Should stop after first call
            ),
            (
                {"WATSONX_API_KEY": "watsonx-api-key"},
                "watsonx-api-key",
                [
                    "WX_API_KEY",
                    "WATSONX_API_KEY",
                ],  # Should check WX_API_KEY first, then WATSONX_API_KEY
            ),
            (
                {"WATSONX_APIKEY": "watsonx-apikey"},
                "watsonx-apikey",
                ["WX_API_KEY", "WATSONX_API_KEY", "WATSONX_APIKEY"],
            ),
            (
                {"WATSONX_ZENAPIKEY": "watsonx-zenapikey"},
                "watsonx-zenapikey",
                [
                    "WX_API_KEY",
                    "WATSONX_API_KEY",
                    "WATSONX_APIKEY",
                    "WATSONX_ZENAPIKEY",
                ],
            ),
            # Test that higher priority keys take precedence
            (
                {
                    "WX_API_KEY": "wx-key",
                    "WATSONX_ZENAPIKEY": "zen-key",
                },
                "wx-key",
                ["WX_API_KEY"],  # Should stop after first call
            ),
            (
                {
                    "WATSONX_API_KEY": "watsonx-api-key",
                    "WATSONX_ZENAPIKEY": "zen-key",
                },
                "watsonx-api-key",
                ["WX_API_KEY", "WATSONX_API_KEY"],  # Should stop after WATSONX_API_KEY
            ),
            (
                {
                    "WATSONX_APIKEY": "watsonx-apikey",
                    "WATSONX_ZENAPIKEY": "zen-key",
                },
                "watsonx-apikey",
                ["WX_API_KEY", "WATSONX_API_KEY", "WATSONX_APIKEY"],
            ),
        ]

        for env_keys, expected_key, expected_calls in test_cases:
            mock_get_secret_str.reset_mock()
            mock_client.reset_mock()
            mock_cache.reset_mock()

            # Configure mock to return values based on env_keys
            def get_secret_side_effect(key):
                return env_keys.get(key)

            mock_get_secret_str.side_effect = get_secret_side_effect

            mock_response = MagicMock()
            mock_response.json.return_value = {
                "access_token": "test-token",
                "expires_in": 3600,
            }
            mock_response.raise_for_status = MagicMock()
            mock_client.post.return_value = mock_response

            # Call function
            result = generate_iam_token()

            # Verify the correct key was used
            call_kwargs = mock_client.post.call_args
            assert (
                call_kwargs.kwargs["data"]["apikey"] == expected_key
            ), f"Expected {expected_key} but got {call_kwargs.kwargs['data']['apikey']} for env_keys: {env_keys}"

            # Verify get_secret_str was called with expected keys (checking short-circuit behavior)
            # Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL"), so we filter that out
            actual_calls = [
                call[0][0]
                for call in mock_get_secret_str.call_args_list
                if call[0][0] != "WATSONX_IAM_URL"
            ]
            assert (
                actual_calls == expected_calls
            ), f"Expected calls {expected_calls} but got {actual_calls} for env_keys: {env_keys}"

    @patch("litellm.llms.watsonx.common_utils.iam_token_cache")
    @patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
    @patch("litellm.llms.watsonx.common_utils.get_secret_str")
    def test_generate_iam_token_with_direct_api_key(
        self, mock_get_secret_str, mock_client, mock_cache
    ):
        """Test that when api_key is passed directly, it's used instead of environment variables."""
        # Setup mocks
        mock_cache.get_cache.return_value = None  # Cache miss
        mock_get_secret_str.return_value = "env-key-should-not-be-used"

        mock_response = MagicMock()
        mock_response.json.return_value = {
            "access_token": "test-token-12345",
            "expires_in": 3600,
        }
        mock_response.raise_for_status = MagicMock()
        mock_client.post.return_value = mock_response

        # Call function with direct api_key
        direct_key = "direct-api-key-12345"
        result = generate_iam_token(api_key=direct_key)

        # Verify get_secret_str was NOT called for API keys (since api_key was provided)
        # Note: get_watsonx_iam_url() calls get_secret_str("WATSONX_IAM_URL"), which is expected
        api_key_calls = [
            call[0][0]
            for call in mock_get_secret_str.call_args_list
            if call[0][0] not in ["WATSONX_IAM_URL"]
        ]
        assert (
            len(api_key_calls) == 0
        ), f"Expected no API key calls but got {api_key_calls}"

        # Verify the direct key was used
        assert result == "test-token-12345"
        call_kwargs = mock_client.post.call_args
        assert call_kwargs.kwargs["data"]["apikey"] == direct_key

    @patch("litellm.llms.watsonx.common_utils.iam_token_cache")
    @patch("litellm.llms.watsonx.common_utils.get_secret_str")
    def test_generate_iam_token_no_api_key_raises_error(
        self, mock_get_secret_str, mock_cache
    ):
        """Test that ValueError is raised when no API key is available."""
        # Setup mocks
        mock_cache.get_cache.return_value = None  # Cache miss
        mock_get_secret_str.return_value = None  # No keys available

        # Call function without api_key and expect ValueError
        with pytest.raises(ValueError, match="API key is required"):
            generate_iam_token()

        # Verify get_secret_str was called for all possible API keys
        # Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL")
        calls = [
            call[0][0]
            for call in mock_get_secret_str.call_args_list
            if call[0][0] != "WATSONX_IAM_URL"
        ]
        assert "WX_API_KEY" in calls
        assert "WATSONX_API_KEY" in calls
        assert "WATSONX_APIKEY" in calls
        assert "WATSONX_ZENAPIKEY" in calls

    @patch("litellm.llms.watsonx.common_utils.iam_token_cache")
    @patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
    @patch("litellm.llms.watsonx.common_utils.get_secret_str")
    def test_generate_iam_token_uses_cache(
        self, mock_get_secret_str, mock_client, mock_cache
    ):
        """Test that cached token is returned when available."""
        # Setup mocks
        cached_token = "cached-token-12345"
        mock_cache.get_cache.return_value = cached_token

        # Call function
        result = generate_iam_token()

        # Verify cached token was returned
        assert result == cached_token

        # Verify get_secret_str and client.post were NOT called (cache hit)
        mock_get_secret_str.assert_not_called()
        mock_client.post.assert_not_called()
