# tests/llm_translation/test_base_aws_llm.py
import os
import json
import pytest
from unittest.mock import patch
from botocore.credentials import Credentials
import sys

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

import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from unittest.mock import Mock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM

import json
import pytest
from unittest.mock import patch, Mock

import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM


def test_bedrock_completion_with_region_name():
    litellm._turn_on_debug()
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        # Construct a response similar to our other tests.
        mock_response.text = json.dumps(
            {
                "response_id": "379ed018/60744aff-e741-4aad-bd10-74639a4ade79",
                "text": "Hello! How's it going? I hope you're having a fantastic day!",
                "generation_id": "38709bb9-f20f-42d9-9c61-13a73b7bbc12",
                "chat_history": [
                    {"role": "USER", "message": "Hello, world!"},
                    {
                        "role": "CHATBOT",
                        "message": "Hello! How's it going? I hope you're having a fantastic day!",
                    },
                ],
                "finish_reason": "COMPLETE",
            }
        )
        mock_response.status_code = 200
        mock_response.headers = {"Content-Type": "application/json"}
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Pass the client so that the HTTP call will be intercepted.
        response = litellm.completion(
            model="cohere.command-r-v1:0",
            messages=[{"role": "user", "content": "Hello, world!"}],
            aws_region_name="us-west-12",
            client=client,
        )

        # Ensure our post method has been called.
        mock_post.assert_called_once()

        assert (
            mock_post.call_args.kwargs["url"]
            == "https://bedrock-runtime.us-west-12.amazonaws.com/model/cohere.command-r-v1:0/invoke"
        )
        assert (
            mock_post.call_args.kwargs["data"]
            == '{"message": "Hello, world!", "chat_history": []}'
        )

        # Print the URL and body of the HTTP request.
        # assert request was signed with the correct region
        _authorization_header = mock_post.call_args.kwargs["headers"]["Authorization"]
        import re

        # Ensure the authorization header contains the exact region segment "us-west-12/bedrock/aws4_request"
        pattern = r"us-west-12/bedrock/aws4_request"
        assert re.search(pattern, _authorization_header) is not None


def test_bedrock_completion_with_dynamic_authentication_params():
    litellm._turn_on_debug()
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        # Construct a response similar to our other tests.
        mock_response.text = json.dumps(
            {
                "response_id": "379ed018/60744aff-e741-4aad-bd10-74639a4ade79",
                "text": "Hello! How's it going? I hope you're having a fantastic day!",
                "generation_id": "38709bb9-f20f-42d9-9c61-13a73b7bbc12",
                "chat_history": [
                    {"role": "USER", "message": "Hello, world!"},
                    {
                        "role": "CHATBOT",
                        "message": "Hello! How's it going? I hope you're having a fantastic day!",
                    },
                ],
                "finish_reason": "COMPLETE",
            }
        )
        mock_response.status_code = 200
        mock_response.headers = {"Content-Type": "application/json"}
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Pass the client so that the HTTP call will be intercepted.
        response = litellm.completion(
            model="cohere.command-r-v1:0",
            messages=[{"role": "user", "content": "Hello, world!"}],
            aws_access_key_id="dynamically_generated_access_key_id",
            aws_secret_access_key="dynamically_generated_secret_access_key",
            client=client,
        )

        # Ensure our post method has been called.
        mock_post.assert_called_once()
        import re

        # Get authorization header
        _authorization_header = mock_post.call_args.kwargs["headers"]["Authorization"]

        # Check for exact credential pattern
        pattern = r"AWS4-HMAC-SHA256 Credential=dynamically_generated_access_key_id/\d{8}/[a-z0-9-]+/bedrock/aws4_request"
        assert re.search(pattern, _authorization_header) is not None


def test_bedrock_completion_with_dynamic_bedrock_runtime_endpoint():
    litellm._turn_on_debug()
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        # Construct a response similar to our other tests.
        mock_response.text = json.dumps(
            {
                "response_id": "379ed018/60744aff-e741-4aad-bd10-74639a4ade79",
                "text": "Hello! How's it going? I hope you're having a fantastic day!",
                "generation_id": "38709bb9-f20f-42d9-9c61-13a73b7bbc12",
                "chat_history": [
                    {"role": "USER", "message": "Hello, world!"},
                    {
                        "role": "CHATBOT",
                        "message": "Hello! How's it going? I hope you're having a fantastic day!",
                    },
                ],
                "finish_reason": "COMPLETE",
            }
        )
        mock_response.status_code = 200
        mock_response.headers = {"Content-Type": "application/json"}
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Pass the client so that the HTTP call will be intercepted.
        response = litellm.completion(
            model="cohere.command-r-v1:0",
            messages=[{"role": "user", "content": "Hello, world!"}],
            aws_bedrock_runtime_endpoint="https://my-fake-endpoint.com",
            client=client,
        )

        # Ensure our post method has been called.
        mock_post.assert_called_once()
        assert (
            mock_post.call_args.kwargs["url"]
            == "https://my-fake-endpoint.com/model/cohere.command-r-v1:0/invoke"
        )


# ------------------------------------------------------------------------------
# A dummy credentials object to return from get_credentials.
# (It must have attributes so that SigV4Auth.add_auth doesn't break.)
# ------------------------------------------------------------------------------
class DummyCredentials:
    access_key = "dummy_access"
    secret_key = "dummy_secret"
    token = "dummy_token"


# ------------------------------------------------------------------------------
# This test makes sure that a given dynamic parameter is passed into the call
# to BaseAWSLLM.get_credentials. (Some dynamic params—for example aws_region_name
# or aws_bedrock_runtime_endpoint—are already covered by other tests.)
# ------------------------------------------------------------------------------
@pytest.mark.parametrize(
    "model",
    [
        "bedrock/converse/cohere.command-r-v1:0",
        "cohere.command-r-v1:0",
        "bedrock/cohere.command-r-v1:0",
        "bedrock/invoke/cohere.command-r-v1:0",
    ],
)
@pytest.mark.parametrize(
    "param_name, param_value",
    [
        ("aws_session_token", "dummy_session_token"),
        ("aws_session_name", "dummy_session_name"),
        ("aws_profile_name", "dummy_profile_name"),
        ("aws_role_name", "dummy_role_name"),
        ("aws_web_identity_token", "dummy_web_identity_token"),
        ("aws_sts_endpoint", "dummy_sts_endpoint"),
    ],
)
def test_dynamic_aws_params_propagation(model, param_name, param_value):
    """
    When passed to litellm.completion, each dynamic AWS authentication parameter
    should propagate down to the get_credentials() call in BaseAWSLLM.

    Also tests different model parameter values.
    """
    client = HTTPHandler()

    # Base parameters required for the completion call.
    # (We include aws_access_key_id and aws_secret_access_key so that the correct auth
    # branch in get_credentials() is reached.)
    base_params = {
        "model": model,
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "aws_access_key_id": "dummy_access",
        "aws_secret_access_key": "dummy_secret",
        "client": client,
    }
    # For parameters such as aws_role_name or aws_web_identity_token a session name is required.
    if param_name in ("aws_role_name", "aws_web_identity_token"):
        base_params["aws_session_name"] = "dummy_session_name"
        if param_name == "aws_web_identity_token":
            # The web identity branch also requires a role name.
            base_params["aws_role_name"] = "dummy_role_name"
    # Inject the dynamic parameter under test.
    base_params[param_name] = param_value

    # Patch SigV4Auth in the signing (so that no actual signing is done).
    with patch("botocore.auth.SigV4Auth", autospec=True) as mock_sigv4:
        instance = mock_sigv4.return_value
        instance.add_auth.return_value = None

        # Patch BaseAWSLLM.get_credentials so that we can capture its kwargs.
        def dummy_get_credentials(**kwargs):
            dummy_get_credentials.called_kwargs = kwargs  # type: ignore[attr-defined]
            return DummyCredentials()

        with patch.object(
            BaseAWSLLM, "get_credentials", side_effect=dummy_get_credentials
        ):
            # Patch the HTTP client's post method to avoid an actual HTTP call.
            with patch.object(client, "post") as mock_post:
                mock_response = Mock()
                mock_response.text = json.dumps(
                    {
                        "response_id": "dummy_response",
                        "text": "Hello! world",
                        "generation_id": "dummy_gen",
                        "chat_history": [],
                        "finish_reason": "COMPLETE",
                    }
                )
                if "converse" in model:
                    mock_response.text = json.dumps(
                        {
                            "output": {
                                "message": {
                                    "role": "assistant",
                                    "content": [{"text": "Here's a joke..."}],
                                }
                            },
                            "usage": {
                                "inputTokens": 12,
                                "outputTokens": 6,
                                "totalTokens": 18,
                            },
                            "stopReason": "stop",
                        }
                    )

                mock_response.status_code = 200
                mock_response.headers = {"Content-Type": "application/json"}
                mock_response.json = lambda: json.loads(mock_response.text)
                mock_post.return_value = mock_response

                # Call litellm.completion with our base & dynamic parameters.
                litellm.completion(**base_params)

                print(
                    "get_credentials.called_kwargs",
                    json.dumps(dummy_get_credentials.called_kwargs, indent=4),
                )

                # We now assert that get_credentials() was called with the dynamic param.
                assert (
                    dummy_get_credentials.called_kwargs.get(param_name) == param_value
                )
