import json
import os
import sys

import pytest
from fastapi.testclient import TestClient

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


from datetime import datetime, timedelta, timezone
from typing import Any, Dict
from unittest.mock import MagicMock, patch

from botocore.awsrequest import AWSPreparedRequest, AWSRequest
from botocore.credentials import Credentials

import litellm
from litellm.llms.bedrock.base_aws_llm import (
    AwsAuthError,
    BaseAWSLLM,
    Boto3CredentialsInfo,
)

# Global variable for the base_aws_llm.py file path

BASE_AWS_LLM_PATH = os.path.join(
    os.path.dirname(__file__), "../../../../litellm/llms/bedrock/base_aws_llm.py"
)


@pytest.fixture(autouse=True)
def flush_shared_bedrock_iam_cache():
    """Process-wide IAM cache must not leak static/env credential entries across tests."""
    BaseAWSLLM._shared_iam_cache.flush_cache()
    yield


def test_base_aws_llm_instances_share_process_wide_iam_cache():
    """Regression LIT-2662: new instances must reuse iam_cache (Bedrock passthrough is per-request)."""
    first = BaseAWSLLM()
    second = BaseAWSLLM()
    assert first.iam_cache is second.iam_cache
    assert first.iam_cache is BaseAWSLLM._shared_iam_cache


def test_static_access_key_credentials_use_iam_cache_across_calls():
    """Static access-key path hits shared iam_cache; second identical call does not refetch."""
    base = BaseAWSLLM()
    fake_creds = MagicMock()

    with patch.object(
        base,
        "_auth_with_access_key_and_secret_key",
        return_value=(fake_creds, 3600),
    ) as mock_static_auth:
        base.get_credentials(
            aws_access_key_id="AKIAEXAMPLE",
            aws_secret_access_key="secret",
            aws_region_name="us-east-1",
        )
        base.get_credentials(
            aws_access_key_id="AKIAEXAMPLE",
            aws_secret_access_key="secret",
            aws_region_name="us-east-1",
        )
        mock_static_auth.assert_called_once()


def _os_environ_without_aws_keys() -> Dict[str, str]:
    """Strip AWS_* so get_credentials hits the ambient-env branch when no explicit keys are passed."""
    return {k: v for k, v in os.environ.items() if not k.startswith("AWS_")}


def test_ambient_env_credentials_use_iam_cache_across_instances():
    """Else-branch env path uses shared iam_cache; second call on another instance does not refetch."""
    base_a = BaseAWSLLM()
    base_b = BaseAWSLLM()
    fake_creds = MagicMock()
    with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True):
        with patch.object(
            BaseAWSLLM,
            "_auth_with_env_vars",
            return_value=(fake_creds, None),
        ) as mock_env:
            base_a.get_credentials()
            base_b.get_credentials()
            mock_env.assert_called_once()


def test_static_access_key_path_boto3_session_constructed_once_when_cached():
    """With real _auth_with_access_key_and_secret_key, boto3.Session is only built once per cache key."""
    base_a = BaseAWSLLM()
    base_b = BaseAWSLLM()
    real_creds = Credentials("AKIAEXAMPLE", "secret-key-val", None)
    mock_session_instance = MagicMock()
    mock_session_instance.get_credentials.return_value = real_creds
    with patch("boto3.Session", return_value=mock_session_instance) as mock_session_cls:
        base_a.get_credentials(
            aws_access_key_id="AKIAEXAMPLE",
            aws_secret_access_key="secret-key-val",
            aws_region_name="us-east-1",
        )
        base_b.get_credentials(
            aws_access_key_id="AKIAEXAMPLE",
            aws_secret_access_key="secret-key-val",
            aws_region_name="us-east-1",
        )
        mock_session_cls.assert_called_once()


def test_ambient_env_path_boto3_session_constructed_once_when_cached():
    """Else branch: boto3.Session() inside _auth_with_env_vars runs once for two cache hits."""
    base_a = BaseAWSLLM()
    base_b = BaseAWSLLM()
    real_creds = Credentials("AKIAENV", "secret-env", None)
    mock_session_instance = MagicMock()
    mock_session_instance.get_credentials.return_value = real_creds
    with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True):
        with patch(
            "boto3.Session", return_value=mock_session_instance
        ) as mock_session_cls:
            base_a.get_credentials()
            base_b.get_credentials()
            mock_session_cls.assert_called_once()


def test_explicit_session_token_tuple_not_cached_in_iam_cache():
    """Temporary key+secret+session paths must not use process-wide iam_cache between calls."""
    base = BaseAWSLLM()
    with patch.object(
        base,
        "_auth_with_aws_session_token",
        return_value=(
            Credentials("ak", "sk", "token"),
            None,
        ),
    ) as mock_sess:
        base.get_credentials(
            aws_access_key_id="AKIA",
            aws_secret_access_key="sec",
            aws_session_token="tok",
        )
        base.get_credentials(
            aws_access_key_id="AKIA",
            aws_secret_access_key="sec",
            aws_session_token="tok",
        )
        assert mock_sess.call_count == 2


def test_aws_profile_path_not_cached_in_iam_cache():
    base = BaseAWSLLM()
    with patch.object(
        base,
        "_auth_with_aws_profile",
        return_value=(Credentials("prof-ak", "prof-sk", None), None),
    ) as mock_profile:
        base.get_credentials(aws_profile_name="my-profile")
        base.get_credentials(aws_profile_name="my-profile")
        assert mock_profile.call_count == 2


def test_web_identity_path_not_cached_in_iam_cache():
    base = BaseAWSLLM()
    with patch.object(
        base,
        "_auth_with_web_identity_token",
        return_value=(Credentials("wi-ak", "wi-sk", "wi-tok"), None),
    ) as mock_wi:
        kwargs = dict(
            aws_web_identity_token="jwt-token",
            aws_role_name="arn:aws:iam::123456789012:role/WebIdentity",
            aws_session_name="web-id-session",
        )
        base.get_credentials(**kwargs)
        base.get_credentials(**kwargs)
        assert mock_wi.call_count == 2


def test_boto3_init_tracer_wrapping():
    """
    Test that all boto3 initializations are wrapped in tracer.trace or @tracer.wrap

    Ensures observability of boto3 calls in litellm.
    """
    # Get the source code of base_aws_llm.py
    with open(BASE_AWS_LLM_PATH, "r") as f:
        content = f.read()

    # List all boto3 initialization patterns we want to check
    boto3_init_patterns = ["boto3.client", "boto3.Session"]

    lines = content.split("\n")
    # Check each boto3 initialization is wrapped in tracer.trace
    for line_number, line in enumerate(lines, 1):
        if line.lstrip().startswith("#"):
            continue
        for pattern in boto3_init_patterns:
            if pattern in line:
                # Look back up to 5 lines for decorator or trace block
                start_line = max(0, line_number - 5)
                context_lines = lines[start_line:line_number]

                has_trace = (
                    "tracer.trace" in line
                    or any("tracer.trace" in prev_line for prev_line in context_lines)
                    or any("@tracer.wrap" in prev_line for prev_line in context_lines)
                )

                if not has_trace:
                    print(f"\nContext for line {line_number}:")
                    for i, ctx_line in enumerate(context_lines, start=start_line + 1):
                        print(f"{i}: {ctx_line}")

                assert (
                    has_trace
                ), f"boto3 initialization '{pattern}' on line {line_number} is not wrapped with tracer.trace or @tracer.wrap"


def test_auth_functions_tracer_wrapping():
    """
    Test that all _auth functions in base_aws_llm.py are wrapped with @tracer.wrap

    Ensures observability of AWS authentication calls in litellm.
    """
    # Get the source code of base_aws_llm.py
    with open(BASE_AWS_LLM_PATH, "r") as f:
        content = f.read()

    lines = content.split("\n")
    # Check each line for _auth function definitions
    for line_number, line in enumerate(lines, 1):
        if line.strip().startswith("def _auth_"):
            # Look back up to 2 lines for the @tracer.wrap decorator
            start_line = max(0, line_number - 2)
            context_lines = lines[start_line:line_number]

            has_tracer_wrap = any(
                "@tracer.wrap" in prev_line for prev_line in context_lines
            )

            if not has_tracer_wrap:
                print(f"\nContext for line {line_number}:")
                for i, ctx_line in enumerate(context_lines, start=start_line + 1):
                    print(f"{i}: {ctx_line}")

            assert (
                has_tracer_wrap
            ), f"Auth function on line {line_number} is not wrapped with @tracer.wrap: {line.strip()}"


def test_get_aws_region_name_boto3_fallback():
    """
    Test the boto3 session fallback logic in _get_aws_region_name method.

    This tests the specific code block that tries to get the region from boto3.Session()
    when aws_region_name is None and not found in environment variables.
    """
    base_aws_llm = BaseAWSLLM()

    # Test case 1: boto3.Session() returns a configured region
    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:
        mock_get_secret.return_value = None  # No region in env vars

        with patch("boto3.Session") as mock_boto3_session:
            mock_session = MagicMock()
            mock_session.region_name = "us-east-1"
            mock_boto3_session.return_value = mock_session

            optional_params = {}
            result = base_aws_llm._get_aws_region_name(optional_params)

            assert result == "us-east-1"
            mock_boto3_session.assert_called_once()

    # Test case 2: boto3.Session() returns None for region (should default to us-west-2)
    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:
        mock_get_secret.return_value = None  # No region in env vars

        with patch("boto3.Session") as mock_boto3_session:
            mock_session = MagicMock()
            mock_session.region_name = None
            mock_boto3_session.return_value = mock_session

            optional_params = {}
            result = base_aws_llm._get_aws_region_name(optional_params)

            assert result == "us-west-2"
            mock_boto3_session.assert_called_once()

    # Test case 3: boto3 import/session creation raises exception (should default to us-west-2)
    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:
        mock_get_secret.return_value = None  # No region in env vars

        with patch("boto3.Session") as mock_boto3_session:
            mock_boto3_session.side_effect = Exception("boto3 not available")

            optional_params = {}
            result = base_aws_llm._get_aws_region_name(optional_params)

            assert result == "us-west-2"
            mock_boto3_session.assert_called_once()

    # Test case 4: aws_region_name is provided in optional_params (should not use boto3)
    with patch("boto3.Session") as mock_boto3_session:
        optional_params = {"aws_region_name": "eu-west-1"}
        result = base_aws_llm._get_aws_region_name(optional_params)

        assert result == "eu-west-1"
        mock_boto3_session.assert_not_called()

    # Test case 5: aws_region_name found in environment variables (should not use boto3)
    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:

        def side_effect(key, default=None):
            if key == "AWS_REGION_NAME":
                return "ap-southeast-1"
            return default

        mock_get_secret.side_effect = side_effect

        with patch("boto3.Session") as mock_boto3_session:
            optional_params = {}
            result = base_aws_llm._get_aws_region_name(optional_params)

            assert result == "ap-southeast-1"
            mock_boto3_session.assert_not_called()


@pytest.mark.parametrize(
    "bad_region",
    [
        "us-east-1@example.com/",
        "us-east-1@example.com",
        "us-east-1/path",
        "us-east-1.example.com",
        "us-east-1:8080",
        "us-east-1#fragment",
        "us-east-1?query=1",
        "us-east-1\\path",
        "US-EAST-1",  # uppercase not allowed
        "us east 1",  # spaces not allowed
        "",  # empty string not allowed
        "us-east-1\n",  # trailing newline must not slip past $
    ],
)
def test_get_aws_region_name_rejects_malformed_region(bad_region):
    """
    Region names are interpolated into endpoint URL templates, so any value
    containing characters that would alter URL parsing must be rejected.
    """
    base_aws_llm = BaseAWSLLM()

    with pytest.raises(ValueError, match="Invalid AWS region format"):
        base_aws_llm._get_aws_region_name(
            optional_params={"aws_region_name": bad_region}
        )


@pytest.mark.parametrize(
    "valid_region",
    [
        "us-east-1",
        "eu-west-2",
        "ap-southeast-1",
        "us-gov-west-1",
        "cn-north-1",
        "me-south-1",
    ],
)
def test_get_aws_region_name_accepts_valid_regions(valid_region):
    """Real AWS region formats must continue to work after the format guard."""
    base_aws_llm = BaseAWSLLM()
    result = base_aws_llm._get_aws_region_name(
        optional_params={"aws_region_name": valid_region}
    )
    assert result == valid_region


def test_get_aws_region_name_rejects_malformed_region_from_env():
    """
    A malformed AWS_REGION / AWS_REGION_NAME env value must also be rejected
    before it can flow into a URL template.
    """
    base_aws_llm = BaseAWSLLM()

    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:

        def side_effect(key, default=None):
            if key == "AWS_REGION_NAME":
                return "us-east-1@example.com/"
            return default

        mock_get_secret.side_effect = side_effect

        with pytest.raises(ValueError, match="Invalid AWS region format"):
            base_aws_llm._get_aws_region_name(optional_params={})


def test_get_aws_region_name_for_non_llm_api_calls_rejects_malformed_param():
    """
    The non-LLM helper (used by Guardrails, Vector Stores, etc.) must validate
    a region passed in directly so it can't flow into a URL template.
    """
    base_aws_llm = BaseAWSLLM()

    with pytest.raises(ValueError, match="Invalid AWS region format"):
        base_aws_llm.get_aws_region_name_for_non_llm_api_calls(
            aws_region_name="us-east-1@example.com/"
        )


def test_get_aws_region_name_for_non_llm_api_calls_rejects_malformed_env():
    """
    A malformed AWS_REGION / AWS_REGION_NAME env value must be rejected on the
    non-LLM path too — Guardrails and Vector Stores read the same env vars.
    """
    base_aws_llm = BaseAWSLLM()

    with patch("litellm.llms.bedrock.base_aws_llm.get_secret") as mock_get_secret:

        def side_effect(key, default=None):
            if key == "AWS_REGION_NAME":
                return "us-east-1@example.com/"
            return default

        mock_get_secret.side_effect = side_effect

        with pytest.raises(ValueError, match="Invalid AWS region format"):
            base_aws_llm.get_aws_region_name_for_non_llm_api_calls()


def test_get_aws_region_name_for_non_llm_api_calls_accepts_valid_region():
    """The non-LLM helper still returns valid regions unchanged."""
    base_aws_llm = BaseAWSLLM()
    assert (
        base_aws_llm.get_aws_region_name_for_non_llm_api_calls(
            aws_region_name="us-east-1"
        )
        == "us-east-1"
    )


def test_get_aws_region_from_model_arn_rejects_malformed_region():
    """
    If the region segment of a model ARN does not match the expected format,
    the helper must return None so the caller falls back to env / default.
    """
    base_aws_llm = BaseAWSLLM()

    bad_arn = (
        "arn:aws:bedrock:us-east-1@example.com:123456789012"
        ":foundation-model/anthropic.claude-3-sonnet"
    )
    assert base_aws_llm._get_aws_region_from_model_arn(bad_arn) is None

    good_arn = (
        "arn:aws:bedrock:us-east-1:123456789012"
        ":foundation-model/anthropic.claude-3-sonnet"
    )
    assert base_aws_llm._get_aws_region_from_model_arn(good_arn) == "us-east-1"


def test_sign_request_with_env_var_bearer_token():
    # Create instance of actual class
    llm = BaseAWSLLM()

    # Test data
    service_name = "bedrock"
    headers = {"Custom-Header": "test"}
    optional_params = {}
    request_data = {"prompt": "test"}
    api_base = "https://api.example.com"

    # Mock environment variable
    with patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": "test_token"}):
        # Execute
        result_headers, result_body = llm._sign_request(
            service_name=service_name,
            headers=headers,
            optional_params=optional_params,
            request_data=request_data,
            api_base=api_base,
        )

        # Assert
        assert result_headers["Authorization"] == "Bearer test_token"
        assert result_headers["Content-Type"] == "application/json"
        assert result_headers["Custom-Header"] == "test"
        assert result_body == json.dumps(request_data).encode()


def test_sign_request_with_sigv4():
    llm = BaseAWSLLM()

    # Mock AWS credentials and SigV4 auth
    mock_credentials = Credentials("test_key", "test_secret", "test_token")
    mock_sigv4 = MagicMock()
    mock_request = MagicMock()
    mock_request.headers = {
        "Authorization": "AWS4-HMAC-SHA256 Credential=test",
        "Content-Type": "application/json",
    }
    mock_request.body = b'{"prompt": "test"}'

    # Test data
    service_name = "bedrock"
    headers = {"Custom-Header": "test"}
    optional_params = {
        "aws_access_key_id": "test_key",
        "aws_secret_access_key": "test_secret",
        "aws_region_name": "us-west-2",
    }
    request_data = {"prompt": "test"}
    api_base = "https://api.example.com"

    # Mock the necessary components
    with (
        patch("botocore.auth.SigV4Auth", return_value=mock_sigv4),
        patch("botocore.awsrequest.AWSRequest", return_value=mock_request),
        patch.object(llm, "get_credentials", return_value=mock_credentials),
        patch.object(llm, "_get_aws_region_name", return_value="us-west-2"),
    ):
        result_headers, result_body = llm._sign_request(
            service_name=service_name,
            headers=headers,
            optional_params=optional_params,
            request_data=request_data,
            api_base=api_base,
        )

        # Assert
        assert "Authorization" in result_headers
        assert result_headers["Authorization"] != "Bearer test_token"
        assert result_headers["Content-Type"] == "application/json"
        assert result_body == mock_request.body


def test_sign_request_with_api_key_bearer_token():
    """
    Test that _sign_request uses the api_key parameter as a bearer token when provided
    """
    llm = BaseAWSLLM()

    # Test data
    service_name = "bedrock"
    headers = {"Custom-Header": "test"}
    optional_params = {}
    request_data = {"prompt": "test"}
    api_base = "https://api.example.com"
    api_key = "test_api_key"

    # Execute with api_key parameter
    result_headers, result_body = llm._sign_request(
        service_name=service_name,
        headers=headers,
        optional_params=optional_params,
        request_data=request_data,
        api_base=api_base,
        api_key=api_key,
    )

    # Assert
    assert result_headers["Authorization"] == f"Bearer {api_key}"
    assert result_headers["Content-Type"] == "application/json"
    assert result_headers["Custom-Header"] == "test"
    assert result_body == json.dumps(request_data).encode()


def test_get_request_headers_with_env_var_bearer_token():
    # Setup
    llm = BaseAWSLLM()
    credentials = Credentials("test_key", "test_secret", "test_token")
    headers = {"Content-Type": "application/json"}
    headers_dict = headers.copy()

    # Create mock request
    mock_prepared_request = MagicMock(spec=AWSPreparedRequest)
    mock_request = MagicMock(spec=AWSRequest)
    mock_request.headers = headers_dict
    mock_request.prepare.return_value = mock_prepared_request

    def mock_aws_request_init(method, url, data, headers):
        mock_request.headers.update(headers)
        return mock_request

    # Test with bearer token
    with (
        patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": "test_token"}),
        patch("botocore.awsrequest.AWSRequest", side_effect=mock_aws_request_init),
    ):
        result = llm.get_request_headers(
            credentials=credentials,
            aws_region_name="us-west-2",
            extra_headers=None,
            endpoint_url="https://api.example.com",
            data='{"prompt": "test"}',
            headers=headers_dict,
        )

        # Assert
        assert mock_request.headers["Authorization"] == "Bearer test_token"
        assert result == mock_prepared_request


def test_get_request_headers_with_sigv4():
    # Setup
    llm = BaseAWSLLM()
    credentials = Credentials("test_key", "test_secret", "test_token")
    headers = {"Content-Type": "application/json"}

    # Create mock request and SigV4 instance
    mock_request = MagicMock(spec=AWSRequest)
    mock_request.headers = headers.copy()
    mock_request.prepare.return_value = MagicMock(spec=AWSPreparedRequest)

    mock_sigv4 = MagicMock()

    # Test without bearer token (should use SigV4)
    with (
        patch.dict(os.environ, {}, clear=True),
        patch("botocore.auth.SigV4Auth", return_value=mock_sigv4) as mock_sigv4_class,
        patch("botocore.awsrequest.AWSRequest", return_value=mock_request),
    ):
        result = llm.get_request_headers(
            credentials=credentials,
            aws_region_name="us-west-2",
            extra_headers=None,
            endpoint_url="https://api.example.com",
            data='{"prompt": "test"}',
            headers=headers,
        )

        # Verify SigV4 authentication and result
        mock_sigv4_class.assert_called_once_with(credentials, "bedrock", "us-west-2")
        mock_sigv4.add_auth.assert_called_once_with(mock_request)
        assert result == mock_request.prepare.return_value


def test_get_request_headers_with_api_key_bearer_token():
    """
    Test that get_request_headers uses the api_key parameter as a bearer token when provided
    """
    # Setup
    llm = BaseAWSLLM()
    credentials = Credentials("test_key", "test_secret", "test_token")
    headers = {"Content-Type": "application/json"}
    headers_dict = headers.copy()
    api_key = "test_api_key"

    # Create mock request
    mock_prepared_request = MagicMock(spec=AWSPreparedRequest)
    mock_request = MagicMock(spec=AWSRequest)
    mock_request.headers = headers_dict
    mock_request.prepare.return_value = mock_prepared_request

    def mock_aws_request_init(method, url, data, headers):
        mock_request.headers.update(headers)
        return mock_request

    # Test with api_key parameter
    with (
        patch.dict(os.environ, {}, clear=True),
        patch("botocore.awsrequest.AWSRequest", side_effect=mock_aws_request_init),
    ):
        result = llm.get_request_headers(
            credentials=credentials,
            aws_region_name="us-west-2",
            extra_headers=None,
            endpoint_url="https://api.example.com",
            data='{"prompt": "test"}',
            headers=headers_dict,
            api_key=api_key,
        )

        # Assert
        assert mock_request.headers["Authorization"] == f"Bearer {api_key}"
        assert result == mock_prepared_request


def test_role_assumption_without_session_name():
    """
    Test for issue 12583: Role assumption should work when only aws_role_name is provided
    without aws_session_name. The system should auto-generate a session name.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Mock the STS response with proper expiration handling
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    # Create a timedelta object that returns 3600 when total_seconds() is called
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    # Test case 1: aws_role_name provided without aws_session_name
    with patch("boto3.client", return_value=mock_sts_client):
        credentials = base_aws_llm.get_credentials(
            aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
        )

        # Verify assume_role was called
        mock_sts_client.assume_role.assert_called_once()

        # Check the call arguments
        call_args = mock_sts_client.assume_role.call_args
        assert (
            call_args[1]["RoleArn"]
            == "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
        )
        # Session name should be auto-generated with format "litellm-session-{timestamp}"
        assert call_args[1]["RoleSessionName"].startswith("litellm-session-")

        # Verify credentials are returned correctly
        assert isinstance(credentials, Credentials)
        assert credentials.access_key == "assumed-access-key"
        assert credentials.secret_key == "assumed-secret-key"
        assert credentials.token == "assumed-session-token"

    # Test case 2: Both aws_role_name and aws_session_name provided (existing behavior)
    mock_sts_client.reset_mock()
    with patch("boto3.client", return_value=mock_sts_client):
        credentials = base_aws_llm.get_credentials(
            aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
            aws_session_name="my-custom-session",
        )

        # Verify assume_role was called with custom session name
        mock_sts_client.assume_role.assert_called_once()
        call_args = mock_sts_client.assume_role.call_args
        assert call_args[1]["RoleSessionName"] == "my-custom-session"

    # Test case 3: AssumeRole is not stored in iam_cache; identical calls each invoke STS.
    BaseAWSLLM._shared_iam_cache.flush_cache()
    mock_sts_client.reset_mock()
    with patch("boto3.client", return_value=mock_sts_client):
        credentials1 = base_aws_llm.get_credentials(
            aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
        )

        credentials2 = base_aws_llm.get_credentials(
            aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
        )

        assert mock_sts_client.assume_role.call_count == 2
        assert credentials1.access_key == credentials2.access_key


def test_assume_role_path_does_not_use_process_iam_cache():
    """AssumeRole credentials are not cached; each get_credentials repeats STS AssumeRole."""
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    mock_sts_client.assume_role.return_value = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::111111111111:assumed-role/SomeOtherRole/session-name",
    }

    role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
    env_without_irsa = {
        k: v
        for k, v in os.environ.items()
        if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
    }

    with patch.dict(os.environ, env_without_irsa, clear=True):
        with patch("boto3.client", return_value=mock_sts_client):
            base_aws_llm.get_credentials(aws_role_name=role_arn)
            mock_sts_client.get_caller_identity.reset_mock()

            base_aws_llm.get_credentials(aws_role_name=role_arn)

            mock_sts_client.get_caller_identity.assert_called()
            assert mock_sts_client.assume_role.call_count == 2


def test_cache_keys_are_different_for_different_roles():
    """
    Test that cache keys are different for different AWS roles.
    This ensures that credentials for different roles don't get mixed up.
    """
    base_aws_llm = BaseAWSLLM()

    # Create arguments for two different roles
    args1 = {
        "aws_access_key_id": None,
        "aws_secret_access_key": None,
        "aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
        "aws_session_name": "test-session-1",
    }

    args2 = {
        "aws_access_key_id": None,
        "aws_secret_access_key": None,
        "aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
        "aws_session_name": "test-session-2",
    }

    # Generate cache keys
    cache_key1 = base_aws_llm.get_cache_key(args1)
    cache_key2 = base_aws_llm.get_cache_key(args2)

    # Cache keys should be different because the role names are different
    assert cache_key1 != cache_key2


def test_different_roles_without_session_names_should_not_share_cache():
    """
    Test that different roles with auto-generated session names don't share cache.
    This was the original issue where cache keys were the same for different roles.
    """
    base_aws_llm = BaseAWSLLM()

    # Create arguments for two different roles without session names
    args1 = {
        "aws_access_key_id": None,
        "aws_secret_access_key": None,
        "aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
        "aws_session_name": None,
    }

    args2 = {
        "aws_access_key_id": None,
        "aws_secret_access_key": None,
        "aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
        "aws_session_name": None,
    }

    # Generate cache keys
    cache_key1 = base_aws_llm.get_cache_key(args1)
    cache_key2 = base_aws_llm.get_cache_key(args2)

    # Cache keys should be different because the role names are different
    assert cache_key1 != cache_key2


@pytest.mark.parametrize(
    "role_kwargs,expected_client_kwargs",
    [
        ({}, {"verify": True}),
        (
            {"aws_region_name": "us-east-1"},
            {"verify": True},
        ),
        (
            {"aws_sts_endpoint": "https://sts.eu-west-1.amazonaws.com"},
            {
                "endpoint_url": "https://sts.eu-west-1.amazonaws.com",
                "region_name": "eu-west-1",
                "verify": True,
            },
        ),
    ],
    ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"],
)
def test_eks_irsa_ambient_credentials_used(role_kwargs, expected_client_kwargs):
    """
    Test that in EKS/IRSA environments, ambient credentials are used when no explicit keys provided.
    This allows web identity tokens to work automatically.
    """
    # Isolate from ambient AWS_REGION/AWS_DEFAULT_REGION so no_region_or_endpoint is deterministic
    env_without_aws_region = {
        k: v
        for k, v in os.environ.items()
        if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
    }
    base_aws_llm = BaseAWSLLM()
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)
    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch.dict(os.environ, env_without_aws_region, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
            credentials, ttl = base_aws_llm._auth_with_aws_role(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                aws_session_name="test-session",
                **role_kwargs,
            )
            mock_boto3_client.assert_called_once_with("sts", **expected_client_kwargs)
            mock_sts_client.assume_role.assert_called_once_with(
                RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                RoleSessionName="test-session",
            )
            assert credentials.access_key == "assumed-access-key"
            assert ttl is not None


@pytest.mark.parametrize(
    "endpoint,expected_region",
    [
        ("https://sts.eu-west-1.amazonaws.com", "eu-west-1"),
        ("https://sts.us-east-1.amazonaws.com", "us-east-1"),
        ("https://sts-fips.us-east-1.amazonaws.com", "us-east-1"),
        ("https://sts-fips.us-gov-west-1.amazonaws.com", "us-gov-west-1"),
        ("https://sts.us-gov-west-1.amazonaws.com", "us-gov-west-1"),
        ("https://sts.cn-north-1.amazonaws.com.cn", "cn-north-1"),
        (
            "https://vpce-abc123.sts.eu-west-1.vpce.amazonaws.com",
            "eu-west-1",
        ),
        ("https://sts.amazonaws.com", None),
        ("https://invalid.example.com", None),
    ],
)
def test_parse_sts_region_from_endpoint(endpoint, expected_region):
    assert BaseAWSLLM._parse_sts_region_from_endpoint(endpoint) == expected_region


@pytest.mark.parametrize(
    "env,aws_sts_endpoint,expected_region",
    [
        ({}, None, None),
        ({"AWS_REGION": "us-east-1"}, None, "us-east-1"),
        ({"AWS_DEFAULT_REGION": "ap-southeast-1"}, None, "ap-southeast-1"),
        ({}, "https://sts.eu-west-1.amazonaws.com", "eu-west-1"),
        (
            {"AWS_REGION": "us-east-1"},
            "https://sts.eu-west-1.amazonaws.com",
            "eu-west-1",
        ),
        ({}, "https://sts.amazonaws.com", None),
        (
            {},
            "https://vpce-abc.sts.eu-central-1.vpce.amazonaws.com",
            "eu-central-1",
        ),
    ],
    ids=[
        "no_env_no_endpoint",
        "env_region",
        "env_default_region",
        "parsed_from_endpoint",
        "parsed_endpoint_over_env",
        "global_endpoint",
        "vpce_endpoint",
    ],
)
def test_resolve_sts_region(env, aws_sts_endpoint, expected_region):
    with patch.dict(os.environ, env, clear=True):
        assert (
            BaseAWSLLM._resolve_sts_region(aws_sts_endpoint=aws_sts_endpoint)
            == expected_region
        )


@pytest.mark.parametrize(
    "env,aws_sts_endpoint,ssl_verify,expected",
    [
        ({}, None, None, {"verify": True}),
        (
            {"AWS_REGION": "us-east-1"},
            None,
            None,
            {"verify": True, "region_name": "us-east-1"},
        ),
        (
            {},
            "https://sts.eu-west-1.amazonaws.com",
            None,
            {
                "verify": True,
                "endpoint_url": "https://sts.eu-west-1.amazonaws.com",
                "region_name": "eu-west-1",
            },
        ),
        (
            {"AWS_REGION": "us-east-1"},
            "https://sts.eu-west-1.amazonaws.com",
            None,
            {
                "verify": True,
                "endpoint_url": "https://sts.eu-west-1.amazonaws.com",
                "region_name": "eu-west-1",
            },
        ),
        (
            {},
            "https://sts.amazonaws.com",
            None,
            {"verify": True, "endpoint_url": "https://sts.amazonaws.com"},
        ),
        (
            {},
            "https://vpce-abc.sts.eu-central-1.vpce.amazonaws.com",
            None,
            {
                "verify": True,
                "endpoint_url": "https://vpce-abc.sts.eu-central-1.vpce.amazonaws.com",
                "region_name": "eu-central-1",
            },
        ),
        ({}, None, False, {"verify": False}),
        (
            {"AWS_DEFAULT_REGION": "ap-southeast-1"},
            None,
            None,
            {"verify": True, "region_name": "ap-southeast-1"},
        ),
    ],
    ids=[
        "default_verify_only",
        "env_region",
        "endpoint_with_parsed_region",
        "endpoint_parsed_over_env",
        "global_endpoint_no_region",
        "vpce_endpoint",
        "ssl_verify_false",
        "env_default_region",
    ],
)
def test_build_sts_client_kwargs(env, aws_sts_endpoint, ssl_verify, expected):
    base_aws_llm = BaseAWSLLM()
    with patch.dict(os.environ, env, clear=True):
        assert (
            base_aws_llm._build_sts_client_kwargs(
                aws_sts_endpoint=aws_sts_endpoint,
                ssl_verify=ssl_verify,
            )
            == expected
        )


def test_irsa_cross_account_sts_client_uses_resolved_region():
    """IRSA cross-account path must use _build_sts_client_kwargs (env region, not Bedrock)."""
    base_aws_llm = BaseAWSLLM()
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
        f.write("test-web-identity-token")
        token_file = f.name

    try:
        with patch.dict(
            os.environ,
            {
                "AWS_WEB_IDENTITY_TOKEN_FILE": token_file,
                "AWS_ROLE_ARN": "arn:aws:iam::111111111111:role/eks-service-account-role",
                "AWS_REGION": "eu-west-1",
            },
            clear=True,
        ):
            mock_sts_client = MagicMock()
            mock_sts_client.assume_role_with_web_identity.return_value = {
                "Credentials": {
                    "AccessKeyId": "temp-key",
                    "SecretAccessKey": "temp-secret",
                    "SessionToken": "temp-token",
                    "Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
                }
            }
            mock_sts_client.assume_role.return_value = {
                "Credentials": {
                    "AccessKeyId": "assumed-key",
                    "SecretAccessKey": "assumed-secret",
                    "SessionToken": "assumed-token",
                    "Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
                }
            }

            with patch(
                "boto3.client", return_value=mock_sts_client
            ) as mock_boto3_client:
                base_aws_llm._auth_with_aws_role(
                    aws_access_key_id=None,
                    aws_secret_access_key=None,
                    aws_session_token=None,
                    aws_role_name="arn:aws:iam::222222222222:role/target-role",
                    aws_session_name="test-session",
                    aws_region_name="eu-central-1",
                )

                for call in mock_boto3_client.call_args_list:
                    assert call.args == ("sts",)
                    assert call.kwargs["region_name"] == "eu-west-1"
                    assert call.kwargs["verify"] is True
    finally:
        os.unlink(token_file)


def test_web_identity_token_sts_client_uses_build_sts_client_kwargs():
    base_aws_llm = BaseAWSLLM()
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role_with_web_identity.return_value = {
        "Credentials": {
            "AccessKeyId": "key",
            "SecretAccessKey": "secret",
            "SessionToken": "token",
            "Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
        },
        "PackedPolicySize": 0,
    }

    with patch.dict(os.environ, {"AWS_REGION": "eu-west-1"}, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
            with patch(
                "litellm.llms.bedrock.base_aws_llm.get_secret",
                return_value="oidc-token",
            ):
                base_aws_llm._auth_with_web_identity_token(
                    aws_web_identity_token="my-token",
                    aws_role_name="arn:aws:iam::111111111111:role/target",
                    aws_session_name="test-session",
                    aws_region_name="eu-central-1",
                    aws_sts_endpoint="https://sts.eu-west-1.amazonaws.com",
                )

                mock_boto3_client.assert_called_once_with(
                    "sts",
                    verify=True,
                    endpoint_url="https://sts.eu-west-1.amazonaws.com",
                    region_name="eu-west-1",
                )


def test_sts_uses_workload_region_not_bedrock_region():
    """Air-gapped: Bedrock in eu-central-1, STS VPC endpoint in eu-west-1 via AWS_REGION."""
    base_aws_llm = BaseAWSLLM()
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }

    with patch.dict(os.environ, {"AWS_REGION": "eu-west-1"}, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
            base_aws_llm._auth_with_aws_role(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                aws_session_name="test-session",
                aws_region_name="eu-central-1",
            )
            mock_boto3_client.assert_called_with(
                "sts",
                region_name="eu-west-1",
                verify=True,
            )


def test_sts_endpoint_region_matches_bedrock_region_param():
    """aws_sts_endpoint signing region must not follow aws_region_name when they differ."""
    base_aws_llm = BaseAWSLLM()
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }

    env_without_irsa = {
        k: v
        for k, v in os.environ.items()
        if k
        not in (
            "AWS_ROLE_ARN",
            "AWS_WEB_IDENTITY_TOKEN_FILE",
            "AWS_REGION",
            "AWS_DEFAULT_REGION",
        )
    }
    with patch.dict(env_without_irsa, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
            base_aws_llm._auth_with_aws_role(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                aws_session_name="test-session",
                aws_region_name="eu-central-1",
                aws_sts_endpoint="https://sts.eu-west-1.amazonaws.com",
            )
            mock_boto3_client.assert_called_with(
                "sts",
                endpoint_url="https://sts.eu-west-1.amazonaws.com",
                region_name="eu-west-1",
                verify=True,
            )


@pytest.mark.parametrize(
    "role_kwargs,expected_client_kwargs",
    [
        (
            {},
            {
                "aws_access_key_id": "explicit-access-key",
                "aws_secret_access_key": "explicit-secret-key",
                "aws_session_token": "assumed-session-token",
                "verify": True,
            },
        ),
        (
            {"aws_region_name": "us-east-1"},
            {
                "aws_access_key_id": "explicit-access-key",
                "aws_secret_access_key": "explicit-secret-key",
                "aws_session_token": "assumed-session-token",
                "verify": True,
            },
        ),
        (
            {"aws_sts_endpoint": "https://sts.eu-west-1.amazonaws.com"},
            {
                "endpoint_url": "https://sts.eu-west-1.amazonaws.com",
                "region_name": "eu-west-1",
                "aws_access_key_id": "explicit-access-key",
                "aws_secret_access_key": "explicit-secret-key",
                "aws_session_token": "assumed-session-token",
                "verify": True,
            },
        ),
    ],
    ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"],
)
def test_explicit_credentials_used_when_provided(role_kwargs, expected_client_kwargs):
    """
    Test that explicit credentials are used when provided (non-EKS/IRSA scenario).
    """
    # Isolate from ambient AWS_REGION/AWS_DEFAULT_REGION so no_region_or_endpoint is deterministic
    env_without_aws_region = {
        k: v
        for k, v in os.environ.items()
        if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
    }
    base_aws_llm = BaseAWSLLM()
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    # Create a timedelta object that returns 3600 when total_seconds() is called
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)
    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch.dict(os.environ, env_without_aws_region, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
            credentials, ttl = base_aws_llm._auth_with_aws_role(
                aws_access_key_id="explicit-access-key",
                aws_secret_access_key="explicit-secret-key",
                aws_session_token="assumed-session-token",
                aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                aws_session_name="test-session",
                **role_kwargs,
            )
            mock_boto3_client.assert_called_once_with("sts", **expected_client_kwargs)
            mock_sts_client.assume_role.assert_called_once_with(
                RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                RoleSessionName="test-session",
            )
            assert credentials.access_key == "assumed-access-key"
            assert credentials.secret_key == "assumed-secret-key"
            assert credentials.token == "assumed-session-token"
            assert ttl is not None


def test_partial_credentials_still_use_ambient():
    """
    Test that if only one credential is provided, we still use ambient credentials.
    This handles edge cases where configuration might be incomplete.
    """
    env_without_aws_region = {
        k: v
        for k, v in os.environ.items()
        if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
    }
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Mock the STS response
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "assumed-access-key",
            "SecretAccessKey": "assumed-secret-key",
            "SessionToken": "assumed-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch.dict(os.environ, env_without_aws_region, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:

            # Call with only access key (missing secret key)
            credentials, ttl = base_aws_llm._auth_with_aws_role(
                aws_access_key_id="AKIAEXAMPLE",
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                aws_session_name="test-session",
            )

            # Should still pass partial credentials to boto3.client
            mock_boto3_client.assert_called_once_with(
                "sts",
                aws_access_key_id="AKIAEXAMPLE",
                aws_secret_access_key=None,
                aws_session_token=None,
                verify=True,
            )

            # Should still call assume_role
            mock_sts_client.assume_role.assert_called_once_with(
                RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
                RoleSessionName="test-session",
            )


def test_cross_account_role_assumption():
    """
    Test assuming a role in a different AWS account (common in multi-account setups).
    """
    env_without_aws_region = {
        k: v
        for k, v in os.environ.items()
        if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
    }
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Mock the STS response for cross-account role
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "cross-account-access-key",
            "SecretAccessKey": "cross-account-secret-key",
            "SessionToken": "cross-account-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch.dict(os.environ, env_without_aws_region, clear=True):
        with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:

            # Assume role in different account (EKS/IRSA scenario)
            credentials, ttl = base_aws_llm._auth_with_aws_role(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
                aws_session_name="cross-account-session",
            )

            # Should use ambient credentials
            mock_boto3_client.assert_called_once_with("sts", verify=True)

            # Should call assume_role with cross-account role
            mock_sts_client.assume_role.assert_called_once_with(
                RoleArn="arn:aws:iam::999999999999:role/CrossAccountRole",
                RoleSessionName="cross-account-session",
            )

            # Verify cross-account credentials are returned
            assert credentials.access_key == "cross-account-access-key"
            assert credentials.secret_key == "cross-account-secret-key"
            assert credentials.token == "cross-account-session-token"
            assert ttl is not None


def test_role_assumption_with_custom_session_name():
    """
    Test role assumption with a custom session name.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Mock the STS response
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "custom-session-access-key",
            "SecretAccessKey": "custom-session-secret-key",
            "SessionToken": "custom-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch("boto3.client", return_value=mock_sts_client):

        # Use custom session name
        credentials, ttl = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
            aws_session_name="evals-bedrock-session",
        )

        # Should call assume_role with custom session name
        mock_sts_client.assume_role.assert_called_once_with(
            RoleArn="arn:aws:iam::1111111111111:role/LitellmRole",
            RoleSessionName="evals-bedrock-session",
        )

        # Verify credentials are returned
        assert credentials.access_key == "custom-session-access-key"
        assert credentials.secret_key == "custom-session-secret-key"
        assert credentials.token == "custom-session-token"


def test_role_assumption_ttl_calculation():
    """
    Test that TTL is calculated correctly from STS response expiration.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Create a real datetime for expiration (1 hour from now)
    expiration_time = datetime.now(timezone.utc) + timedelta(hours=1)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "ttl-test-access-key",
            "SecretAccessKey": "ttl-test-secret-key",
            "SessionToken": "ttl-test-session-token",
            "Expiration": expiration_time,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch("boto3.client", return_value=mock_sts_client):

        credentials, ttl = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
            aws_session_name="ttl-test-session",
        )

        # TTL should be approximately 3540 seconds (1 hour - 60 second buffer)
        assert ttl is not None
        assert 3500 <= ttl <= 3600  # Allow some variance for test execution time


def test_role_assumption_access_denied_falls_back_when_same_role():
    """
    Test that when AssumeRole fails with AccessDenied AND the caller is confirmed
    to already be running as the target role, we fall back to ambient credentials.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client to raise AccessDenied
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.side_effect = Exception(
        "An error occurred (AccessDenied) when calling the AssumeRole operation: "
        "Roles may not be assumed by root accounts."
    )

    # Mock _auth_with_env_vars to return fallback credentials
    mock_creds = MagicMock()
    mock_creds.access_key = "fallback-access-key"
    mock_creds.secret_key = "fallback-secret-key"

    with patch("boto3.client", return_value=mock_sts_client):
        with patch.object(
            base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None)
        ) as mock_env_auth:
            # _is_already_running_as_role returns True => fallback allowed
            with patch.object(
                base_aws_llm, "_is_already_running_as_role", return_value=True
            ):
                credentials, ttl = base_aws_llm._auth_with_aws_role(
                    aws_access_key_id=None,
                    aws_secret_access_key=None,
                    aws_session_token=None,
                    aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
                    aws_session_name="error-test-session",
                )

                # Should have fallen back to env vars
                mock_env_auth.assert_called_once()
                assert credentials.access_key == "fallback-access-key"


def test_role_assumption_access_denied_raises_when_different_role():
    """
    Test that when AssumeRole fails with AccessDenied but the caller is NOT
    the same role, the error is re-raised (genuine permission failure).
    """
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.side_effect = Exception(
        "An error occurred (AccessDenied) when calling the AssumeRole operation: "
        "User is not authorized to perform sts:AssumeRole"
    )

    with patch("boto3.client", return_value=mock_sts_client):
        # _is_already_running_as_role returns False => do NOT fallback
        with patch.object(
            base_aws_llm, "_is_already_running_as_role", return_value=False
        ):
            with pytest.raises(Exception) as exc_info:
                base_aws_llm._auth_with_aws_role(
                    aws_access_key_id=None,
                    aws_secret_access_key=None,
                    aws_session_token=None,
                    aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
                    aws_session_name="error-test-session",
                )

            assert "AccessDenied" in str(exc_info.value)


def test_role_assumption_non_access_denied_error_propagated():
    """
    Test that non-AccessDenied errors from AssumeRole are still propagated.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client to raise a non-AccessDenied exception
    mock_sts_client = MagicMock()
    mock_sts_client.assume_role.side_effect = Exception(
        "An error occurred (MalformedPolicyDocument) when calling the AssumeRole operation"
    )

    with patch("boto3.client", return_value=mock_sts_client):
        with pytest.raises(Exception) as exc_info:
            base_aws_llm._auth_with_aws_role(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_session_token=None,
                aws_role_name="arn:aws:iam::1111111111111:role/BadPolicyRole",
                aws_session_name="error-test-session",
            )

        assert "MalformedPolicyDocument" in str(exc_info.value)


def test_multiple_role_assumptions_in_sequence():
    """
    Test that multiple role assumptions work correctly in sequence.
    This simulates the scenario where different models use different roles.
    """
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()

    # Mock different responses for different roles
    mock_expiry = MagicMock()
    mock_expiry.tzinfo = timezone.utc
    time_diff = MagicMock()
    time_diff.total_seconds.return_value = 3600
    mock_expiry.__sub__ = MagicMock(return_value=time_diff)

    # First role response
    mock_sts_response1 = {
        "Credentials": {
            "AccessKeyId": "role1-access-key",
            "SecretAccessKey": "role1-secret-key",
            "SessionToken": "role1-session-token",
            "Expiration": mock_expiry,
        }
    }

    # Second role response
    mock_sts_response2 = {
        "Credentials": {
            "AccessKeyId": "role2-access-key",
            "SecretAccessKey": "role2-secret-key",
            "SessionToken": "role2-session-token",
            "Expiration": mock_expiry,
        }
    }

    # Configure mock to return different responses
    mock_sts_client.assume_role.side_effect = [mock_sts_response1, mock_sts_response2]

    with patch("boto3.client", return_value=mock_sts_client):

        # First role assumption
        credentials1, ttl1 = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
            aws_session_name="session-1",
        )

        # Second role assumption
        credentials2, ttl2 = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
            aws_session_name="session-2",
        )

        # Verify both role assumptions were made
        assert mock_sts_client.assume_role.call_count == 2

        # Verify first role credentials
        assert credentials1.access_key == "role1-access-key"
        assert credentials1.secret_key == "role1-secret-key"
        assert credentials1.token == "role1-session-token"

        # Verify second role credentials
        assert credentials2.access_key == "role2-access-key"
        assert credentials2.secret_key == "role2-secret-key"
        assert credentials2.token == "role2-session-token"


def test_auth_with_aws_role_irsa_environment():
    """Test that _auth_with_aws_role detects and uses IRSA environment variables"""
    base_llm = BaseAWSLLM()

    # Create a temporary file to simulate the web identity token
    import tempfile

    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
        f.write("test-web-identity-token")
        token_file = f.name

    try:
        # Set IRSA environment variables
        with patch.dict(
            os.environ,
            {
                "AWS_WEB_IDENTITY_TOKEN_FILE": token_file,
                "AWS_ROLE_ARN": "arn:aws:iam::111111111111:role/eks-service-account-role",
                "AWS_REGION": "us-east-1",
            },
        ):
            # Mock the boto3 STS client
            mock_sts_client = MagicMock()
            mock_assume_web_identity_response = {
                "Credentials": {
                    "AccessKeyId": "irsa-temp-access-key",
                    "SecretAccessKey": "irsa-temp-secret-key",
                    "SessionToken": "irsa-temp-session-token",
                    "Expiration": datetime.now() + timedelta(hours=1),
                }
            }
            mock_assume_role_response = {
                "Credentials": {
                    "AccessKeyId": "irsa-access-key",
                    "SecretAccessKey": "irsa-secret-key",
                    "SessionToken": "irsa-session-token",
                    "Expiration": datetime.now() + timedelta(hours=1),
                }
            }
            mock_sts_client.assume_role_with_web_identity.return_value = (
                mock_assume_web_identity_response
            )
            mock_sts_client.assume_role.return_value = mock_assume_role_response

            with patch(
                "boto3.client", return_value=mock_sts_client
            ) as mock_boto3_client:
                # Call _auth_with_aws_role without explicit credentials
                creds, ttl = base_llm._auth_with_aws_role(
                    aws_access_key_id=None,
                    aws_secret_access_key=None,
                    aws_session_token=None,
                    aws_role_name="arn:aws:iam::222222222222:role/target-role",
                    aws_session_name="test-session",
                )

                # Verify boto3.client was called multiple times
                # First for manual IRSA, then with IRSA credentials
                assert mock_boto3_client.call_count >= 2

                # Verify assume_role_with_web_identity was called
                mock_sts_client.assume_role_with_web_identity.assert_called_once_with(
                    RoleArn="arn:aws:iam::111111111111:role/eks-service-account-role",
                    RoleSessionName="test-session",
                    WebIdentityToken="test-web-identity-token",
                )

                # Verify assume_role was called with correct parameters
                mock_sts_client.assume_role.assert_called_once_with(
                    RoleArn="arn:aws:iam::222222222222:role/target-role",
                    RoleSessionName="test-session",
                )

                # Verify the returned credentials
                assert creds.access_key == "irsa-access-key"
                assert creds.secret_key == "irsa-secret-key"
                assert creds.token == "irsa-session-token"
                assert ttl > 0  # TTL should be positive
    finally:
        # Clean up the temporary file
        os.unlink(token_file)


def test_auth_with_aws_role_same_role_irsa():
    """Test that when IRSA role matches the requested role, we skip assumption"""
    base_llm = BaseAWSLLM()

    # Set IRSA environment variables
    with patch.dict(
        os.environ,
        {
            "AWS_ROLE_ARN": "arn:aws:iam::111111111111:role/LitellmRole",
            "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/eks.amazonaws.com/serviceaccount/token",
        },
    ):
        # Mock the _auth_with_env_vars method
        mock_creds = MagicMock()
        mock_creds.access_key = "irsa-access-key"
        mock_creds.secret_key = "irsa-secret-key"
        mock_creds.token = "irsa-session-token"

        with patch.object(
            base_llm, "_auth_with_env_vars", return_value=(mock_creds, None)
        ) as mock_env_auth:
            # Call get_credentials instead of _auth_with_aws_role directly
            # This tests the full flow
            creds = base_llm.get_credentials(
                aws_access_key_id=None,
                aws_secret_access_key=None,
                aws_role_name="arn:aws:iam::111111111111:role/LitellmRole",  # Same as AWS_ROLE_ARN
                aws_session_name="test-session",
                aws_region_name="us-east-1",
            )

            # Verify it used the env vars auth (no role assumption)
            mock_env_auth.assert_called_once()

            # Verify the returned credentials
            assert creds.access_key == "irsa-access-key"


def test_assume_role_with_external_id():
    """Test that assume_role STS call includes ExternalId parameter when provided"""
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()
    mock_expiry = datetime.now(timezone.utc) + timedelta(hours=1)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "test-access-key",
            "SecretAccessKey": "test-secret-key",
            "SessionToken": "test-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch("boto3.client", return_value=mock_sts_client):
        # Call _auth_with_aws_role with external ID
        credentials, ttl = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::123456789012:role/ExampleRole",
            aws_session_name="test-session",
            aws_external_id="UniqueExternalID123",
        )

        # Verify assume_role was called with ExternalId
        mock_sts_client.assume_role.assert_called_once_with(
            RoleArn="arn:aws:iam::123456789012:role/ExampleRole",
            RoleSessionName="test-session",
            ExternalId="UniqueExternalID123",
        )


def test_assume_role_without_external_id():
    """Test that assume_role STS call excludes ExternalId parameter when not provided"""
    base_aws_llm = BaseAWSLLM()

    # Mock the boto3 STS client
    mock_sts_client = MagicMock()
    mock_expiry = datetime.now(timezone.utc) + timedelta(hours=1)

    mock_sts_response = {
        "Credentials": {
            "AccessKeyId": "test-access-key",
            "SecretAccessKey": "test-secret-key",
            "SessionToken": "test-session-token",
            "Expiration": mock_expiry,
        }
    }
    mock_sts_client.assume_role.return_value = mock_sts_response

    with patch("boto3.client", return_value=mock_sts_client):
        # Call _auth_with_aws_role without external ID
        credentials, ttl = base_aws_llm._auth_with_aws_role(
            aws_access_key_id=None,
            aws_secret_access_key=None,
            aws_session_token=None,
            aws_role_name="arn:aws:iam::123456789012:role/ExampleRole",
            aws_session_name="test-session",
        )

        # Verify assume_role was called without ExternalId
        mock_sts_client.assume_role.assert_called_once_with(
            RoleArn="arn:aws:iam::123456789012:role/ExampleRole",
            RoleSessionName="test-session",
        )


def test_converse_handler_external_id_extraction():
    """Test that BedrockConverseLLM properly extracts and passes aws_external_id parameter"""
    from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM

    converse_llm = BedrockConverseLLM()

    # Mock get_credentials to capture parameters
    def mock_get_credentials(**kwargs):
        mock_get_credentials.called_kwargs = kwargs
        mock_credentials = MagicMock()
        mock_credentials.access_key = "test-access-key"
        mock_credentials.secret_key = "test-secret-key"
        mock_credentials.token = "test-session-token"
        return mock_credentials

    with patch.object(
        converse_llm, "get_credentials", side_effect=mock_get_credentials
    ):
        with patch.object(
            converse_llm, "_get_aws_region_name", return_value="us-west-2"
        ):
            with patch.object(
                converse_llm,
                "get_runtime_endpoint",
                return_value=("https://test", "https://test"),
            ):
                with patch("litellm.AmazonConverseConfig") as mock_config:
                    mock_config.return_value._transform_request.return_value = {
                        "test": "data"
                    }
                    with patch.object(
                        converse_llm, "get_request_headers"
                    ) as mock_headers:
                        mock_headers.return_value = MagicMock()
                        mock_headers.return_value.headers = {"Authorization": "test"}
                        with patch(
                            "litellm.llms.custom_httpx.http_handler._get_httpx_client"
                        ) as mock_client:
                            mock_http_client = MagicMock()
                            mock_response = MagicMock()
                            mock_response.raise_for_status.return_value = None
                            mock_http_client.post.return_value = mock_response
                            mock_client.return_value = mock_http_client

                            # Mock the transform_response method
                            mock_config.return_value._transform_response.return_value = (
                                MagicMock()
                            )

                            # Call completion with aws_external_id in optional_params
                            optional_params = {
                                "aws_role_name": "arn:aws:iam::123456789012:role/ExampleRole",
                                "aws_session_name": "test-session",
                                "aws_external_id": "TestExternalID123",
                            }

                            try:
                                converse_llm.completion(
                                    model="anthropic.claude-3-sonnet-20240229-v1:0",
                                    messages=[{"role": "user", "content": "Hello"}],
                                    api_base=None,
                                    custom_prompt_dict={},
                                    model_response=MagicMock(),
                                    encoding="utf-8",
                                    logging_obj=MagicMock(),
                                    optional_params=optional_params,
                                    acompletion=False,
                                    timeout=None,
                                    litellm_params={},
                                )
                            except Exception:
                                # We expect this to fail due to mocking, but that's OK
                                # We just want to verify the parameter extraction
                                pass

                            # Verify aws_external_id was extracted and passed to get_credentials
                            assert hasattr(mock_get_credentials, "called_kwargs")
                            assert (
                                "aws_external_id" in mock_get_credentials.called_kwargs
                            )
                            assert (
                                mock_get_credentials.called_kwargs["aws_external_id"]
                                == "TestExternalID123"
                            )


def test_is_already_running_as_role_irsa_same_role():
    """Test IRSA fast path: when AWS_ROLE_ARN matches target role."""
    base_aws_llm = BaseAWSLLM()

    with patch.dict(
        os.environ,
        {
            "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole",
            "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token",
        },
    ):
        assert (
            base_aws_llm._is_already_running_as_role(
                "arn:aws:iam::123456789012:role/MyRole"
            )
            is True
        )


def test_is_already_running_as_role_irsa_different_role():
    """Test IRSA fast path: when AWS_ROLE_ARN does NOT match target role."""
    base_aws_llm = BaseAWSLLM()

    with patch.dict(
        os.environ,
        {
            "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole",
            "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token",
        },
    ):
        assert (
            base_aws_llm._is_already_running_as_role(
                "arn:aws:iam::999999999999:role/OtherRole"
            )
            is False
        )


def test_is_already_running_as_role_ecs_task_role():
    """Test ECS/EC2 path: GetCallerIdentity shows assumed-role matching target."""
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        # Ensure no IRSA env vars
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws:iam::123456789012:role/MyEcsTaskRole"
                    )
                    is True
                )


def test_is_already_running_as_role_ecs_different_role():
    """Test ECS/EC2 path: GetCallerIdentity shows a different role."""
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws:iam::999999999999:role/DifferentRole"
                    )
                    is False
                )


def test_is_already_running_as_role_ecs_role_with_path():
    """Test ECS path with role that has a path prefix (e.g., /service-role/MyRole)."""
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                # Role ARN with path
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws:iam::123456789012:role/service-role/MyEcsTaskRole"
                    )
                    is True
                )


def test_is_already_running_as_role_get_caller_identity_fails():
    """Test that when GetCallerIdentity fails, we return False (don't crash)."""
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.side_effect = Exception("No credentials found")

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws:iam::123456789012:role/SomeRole"
                    )
                    is False
                )


def test_get_credentials_ecs_same_role_skips_assume_role():
    """
    End-to-end test: when running on ECS with the same role as aws_role_name,
    get_credentials should use ambient credentials and NOT call AssumeRole.
    """
    base_aws_llm = BaseAWSLLM()

    mock_creds = MagicMock()
    mock_creds.access_key = "ecs-access-key"
    mock_creds.secret_key = "ecs-secret-key"
    mock_creds.token = "ecs-session-token"

    with patch.object(
        base_aws_llm,
        "_is_already_running_as_role",
        return_value=True,
    ) as mock_already_running:
        with patch.object(
            base_aws_llm,
            "_auth_with_env_vars",
            return_value=(mock_creds, None),
        ) as mock_env_auth:
            with patch.object(
                base_aws_llm,
                "_auth_with_aws_role",
            ) as mock_role_auth:
                credentials = base_aws_llm.get_credentials(
                    aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole",
                    aws_region_name="us-east-1",
                )
                base_aws_llm.get_credentials(
                    aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole",
                    aws_region_name="us-east-1",
                )

                # Each get_credentials must check identity first; second call still checks before
                # taking the iam_cache hit (no pre-peek that bypasses _is_already_running_as_role).
                assert mock_already_running.call_count == 2
                # Cached env resolution: second call hits iam_cache, not _auth_with_env_vars again.
                mock_env_auth.assert_called_once()
                mock_role_auth.assert_not_called()
                assert credentials.access_key == "ecs-access-key"


def test_get_credentials_role_second_call_not_same_role_uses_assume_not_env_cache():
    """
    First request: already target role -> env path fills iam_cache.
    Second request (e.g. identity changed): not same role -> AssumeRole path; must not reuse
    cached env resolution from the first call.
    """
    base_aws_llm = BaseAWSLLM()

    env_creds = MagicMock()
    env_creds.access_key = "ambient-key"
    env_creds.secret_key = "ambient-secret"
    env_creds.token = "ambient-token"

    assumed_creds = MagicMock()
    assumed_creds.access_key = "assumed-key"
    assumed_creds.secret_key = "assumed-secret"
    assumed_creds.token = "assumed-token"

    role_arn = "arn:aws:iam::123456789012:role/TargetRole"

    with patch.object(
        base_aws_llm,
        "_is_already_running_as_role",
        side_effect=[True, False],
    ) as mock_already:
        with patch.object(
            base_aws_llm,
            "_auth_with_env_vars",
            return_value=(env_creds, None),
        ) as mock_env_auth:
            with patch.object(
                base_aws_llm,
                "_auth_with_aws_role",
                return_value=(assumed_creds, 3600),
            ) as mock_role_auth:
                first = base_aws_llm.get_credentials(
                    aws_role_name=role_arn,
                    aws_region_name="us-east-1",
                )
                second = base_aws_llm.get_credentials(
                    aws_role_name=role_arn,
                    aws_region_name="us-east-1",
                )

                assert mock_already.call_count == 2
                mock_env_auth.assert_called_once()
                mock_role_auth.assert_called_once()
                assert first.access_key == "ambient-key"
                assert second.access_key == "assumed-key"


def test_parse_arn_account_and_role_name():
    """Test the ARN parser helper for various ARN formats."""
    parse = BaseAWSLLM._parse_arn_account_and_role_name

    # Standard IAM role ARN
    assert parse("arn:aws:iam::123456789012:role/MyRole") == (
        "aws",
        "123456789012",
        "MyRole",
    )

    # IAM role ARN with path
    assert parse("arn:aws:iam::123456789012:role/service-role/MyRole") == (
        "aws",
        "123456789012",
        "MyRole",
    )

    # Assumed-role ARN (from GetCallerIdentity)
    assert parse("arn:aws:sts::123456789012:assumed-role/MyRole/session-id") == (
        "aws",
        "123456789012",
        "MyRole",
    )

    # China partition
    assert parse("arn:aws-cn:iam::123456789012:role/MyRole") == (
        "aws-cn",
        "123456789012",
        "MyRole",
    )

    # GovCloud partition
    assert parse("arn:aws-us-gov:iam::123456789012:role/MyRole") == (
        "aws-us-gov",
        "123456789012",
        "MyRole",
    )

    # Invalid ARNs
    assert parse("not-an-arn") is None
    assert parse("arn:aws:iam::123456789012:user/MyUser") is None
    assert parse("") is None


def test_is_already_running_as_role_cross_account_same_name():
    """
    Test that same role NAME in different accounts does NOT match.
    This is the cross-account false-match prevention.
    """
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    # Caller is in account 111111111111
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::111111111111:assumed-role/MyRole/session-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                # Target is same role name but in account 222222222222
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws:iam::222222222222:role/MyRole"
                    )
                    is False
                )


def test_is_already_running_as_role_cross_partition():
    """
    Test that same role name + account but different partition does NOT match.
    """
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch("boto3.client", return_value=mock_sts_client):
                # Same account and role but aws-cn partition
                assert (
                    base_aws_llm._is_already_running_as_role(
                        "arn:aws-cn:iam::123456789012:role/MyRole"
                    )
                    is False
                )


def test_is_already_running_as_role_invalid_target_arn():
    """
    Test that an unparseable target ARN returns False immediately.
    """
    base_aws_llm = BaseAWSLLM()

    # Should return False without making any API calls
    assert base_aws_llm._is_already_running_as_role("not-a-valid-arn") is False


def test_filter_headers_skips_none_values():
    """
    Test that _filter_headers_for_aws_signature skips headers with None values.

    Reproduces the issue where botocore's SigV4Auth crashes with
    'NoneType' object has no attribute 'split' when a header value is None.
    """
    llm = BaseAWSLLM()

    headers = {
        "Content-Type": "application/json",
        "x-amz-security-token": None,
        "x-amzn-bedrock-kb-session-id": None,
        "host": None,
        "x-amz-date": "20240101T000000Z",
        "x-custom-header": None,
    }

    filtered = llm._filter_headers_for_aws_signature(headers)

    assert filtered["Content-Type"] == "application/json"
    assert filtered["x-amz-date"] == "20240101T000000Z"
    assert "x-amz-security-token" not in filtered
    assert "x-amzn-bedrock-kb-session-id" not in filtered
    assert "host" not in filtered
    # Non-AWS headers are excluded regardless
    assert "x-custom-header" not in filtered


def test_sign_request_with_none_header_values():
    """
    End-to-end test that _sign_request does not crash when headers contain
    None values for x-amz-* keys.

    This reproduces the Bedrock KB GovCloud issue where SigV4 signing failed
    with 'NoneType' object has no attribute 'split'.

    Also verifies that None-valued headers are NOT re-merged into the
    returned headers dict (which would cause downstream HTTP client failures).
    """
    llm = BaseAWSLLM()

    mock_credentials = Credentials("test_key", "test_secret")

    headers_with_nones = {
        "Content-Type": "application/json",
        "x-amzn-trace-id": None,
        "x-forwarded-for": None,
    }

    with (
        patch.object(llm, "get_credentials", return_value=mock_credentials),
        patch.object(llm, "_get_aws_region_name", return_value="us-gov-west-1"),
    ):
        result_headers, result_body = llm._sign_request(
            service_name="bedrock",
            headers=headers_with_nones,
            optional_params={
                "aws_access_key_id": "test_key",
                "aws_secret_access_key": "test_secret",
                "aws_region_name": "us-gov-west-1",
            },
            request_data={"retrievalQuery": {"text": "test query"}},
            api_base="https://bedrock-agent-runtime.us-gov-west-1.amazonaws.com/knowledgebases/KB123/retrieve",
        )

        assert "Authorization" in result_headers
        assert result_body is not None

        # None-valued headers must NOT appear in the returned headers
        for header_name, header_value in result_headers.items():
            assert (
                header_value is not None
            ), f"Header '{header_name}' has None value in returned headers"


def test_is_already_running_as_role_ssl_verify_passed():
    """
    Test that ssl_verify parameter is correctly passed to the STS client.
    """
    base_aws_llm = BaseAWSLLM()

    mock_sts_client = MagicMock()
    mock_sts_client.get_caller_identity.return_value = {
        "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id"
    }

    with patch.dict(os.environ, {}, clear=False):
        env = {
            k: v
            for k, v in os.environ.items()
            if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
        }
        with patch.dict(os.environ, env, clear=True):
            with patch(
                "boto3.client", return_value=mock_sts_client
            ) as mock_boto3_client:
                base_aws_llm._is_already_running_as_role(
                    "arn:aws:iam::123456789012:role/MyRole",
                    ssl_verify="/path/to/ca-bundle.crt",
                )
                mock_boto3_client.assert_called_once_with(
                    "sts", verify="/path/to/ca-bundle.crt"
                )


# ---------------------------------------------------------------------------
# LIT-3274: get_bedrock_model_id must strip "bedrock/" prefix and URL-encode
# ARNs for the invoke path (invoke-with-response-stream).  Without this fix
# the Bedrock API receives a malformed URL, returns a JSON error body, and
# botocore's EventStreamBuffer raises ChecksumMismatch instead of the real
# error.  0x223a7b22 == ':{\"' — the start of a JSON object.
# ---------------------------------------------------------------------------


class TestGetBedrockModelIdArnHandling:
    """Unit tests for get_bedrock_model_id with inference-profile ARNs."""

    ARN = "arn:aws:bedrock:us-east-1:086734376398:inference-profile/global.anthropic.claude-sonnet-4-5-20250929-v1:0"

    def _call(self, model: str, optional_params: dict | None = None) -> str:
        from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM

        provider = BaseAWSLLM.get_bedrock_invoke_provider(model)
        return BaseAWSLLM.get_bedrock_model_id(
            model=model,
            provider=provider,
            optional_params=optional_params or {},
        )

    def test_arn_with_bedrock_prefix_is_stripped_and_encoded(self):
        """bedrock/arn:... must not appear verbatim in the model_id."""
        model_id = self._call(f"bedrock/{self.ARN}")
        assert (
            "bedrock/arn" not in model_id
        ), f"'bedrock/' prefix not stripped; got: {model_id}"
        # Must be URL-encoded (colons → %3A)
        assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}"
        assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}"

    def test_arn_with_compound_bedrock_invoke_prefix_is_fully_stripped_and_encoded(
        self,
    ):
        """bedrock/invoke/arn:... — compound prefix — must be fully stripped.

        The old fix used ``break`` after the first matched prefix, so
        ``bedrock/invoke/arn:...`` would only strip ``bedrock/``, leaving
        ``invoke/arn:...``.  The subsequent ``.replace('invoke/', '')`` call
        then returned the bare unencoded ARN, reproducing the same
        malformed-URL bug the fix aimed to prevent.

        strip_bedrock_routing_prefix() has no break and handles this correctly.
        """
        model_id = self._call(f"bedrock/invoke/{self.ARN}")
        assert (
            "invoke/" not in model_id
        ), f"'invoke/' prefix not stripped; got: {model_id}"
        assert (
            "bedrock/" not in model_id
        ), f"'bedrock/' prefix not stripped; got: {model_id}"
        assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}"
        assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}"

    def test_bare_arn_is_encoded(self):
        """Direct ARN without routing prefix must also be URL-encoded."""
        model_id = self._call(self.ARN)
        assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}"
        assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}"

    def test_arn_url_matches_expected(self):
        """Full URL built from messages config must match expected encoded form."""
        import urllib.parse
        from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
            AmazonAnthropicClaudeMessagesConfig,
        )

        config = AmazonAnthropicClaudeMessagesConfig()
        url = config.get_complete_url(
            api_base=None,
            api_key=None,
            model=f"bedrock/{self.ARN}",
            optional_params={"aws_region_name": "us-east-1"},
            litellm_params={},
            stream=True,
        )
        encoded_arn = urllib.parse.quote(self.ARN, safe="")
        expected = (
            f"https://bedrock-runtime.us-east-1.amazonaws.com"
            f"/model/{encoded_arn}/invoke-with-response-stream"
        )
        assert (
            url == expected
        ), f"URL mismatch:\n  got:      {url}\n  expected: {expected}"

    def test_regular_model_id_unaffected(self):
        """Non-ARN model IDs must continue to work as before."""
        model_id = self._call("anthropic.claude-3-sonnet-20240229-v1:0")
        assert model_id == "anthropic.claude-3-sonnet-20240229-v1:0"

    def test_invoke_prefixed_model_unaffected(self):
        """invoke/ prefix stripping still works after the fix."""
        model_id = self._call("invoke/anthropic.claude-3-sonnet-20240229-v1:0")
        assert model_id == "anthropic.claude-3-sonnet-20240229-v1:0"
