import os
from unittest.mock import patch, MagicMock
from crewai import LLM


class TestBedrockAuthentication:
    """Test AWS Bedrock authentication methods."""

    @patch.dict(os.environ, {
        'AWS_ACCESS_KEY_ID': 'test-key-id',
        'AWS_SECRET_ACCESS_KEY': 'test-secret-key', 
        'AWS_DEFAULT_REGION': 'us-east-1'
    })
    @patch('litellm.completion')
    def test_bedrock_iam_authentication(self, mock_completion):
        """Test Bedrock with IAM role authentication."""
        mock_completion.return_value = MagicMock()
        mock_completion.return_value.choices = [MagicMock()]
        mock_completion.return_value.choices[0].message.content = "Test response"
        
        llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
        result = llm.call("test message")
        
        mock_completion.assert_called_once()
        assert result == "Test response"

    @patch.dict(os.environ, {
        'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
        'AWS_DEFAULT_REGION': 'us-east-1'
    })
    @patch('litellm.completion')
    def test_bedrock_api_key_authentication(self, mock_completion):
        """Test Bedrock with API key authentication."""
        mock_completion.return_value = MagicMock()
        mock_completion.return_value.choices = [MagicMock()]
        mock_completion.return_value.choices[0].message.content = "Test response"
        
        llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
        result = llm.call("test message")
        
        mock_completion.assert_called_once()
        assert result == "Test response"

    def test_bedrock_missing_credentials(self):
        """Test Bedrock fails gracefully with missing credentials."""
        with patch.dict(os.environ, {}, clear=True):
            llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
            assert llm.model == "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"

    @patch.dict(os.environ, {
        'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
        'AWS_DEFAULT_REGION': 'us-east-1'
    })
    @patch('litellm.completion')
    def test_bedrock_api_key_with_streaming(self, mock_completion):
        """Test Bedrock API key authentication with streaming."""
        mock_completion.return_value = iter([
            MagicMock(choices=[MagicMock(delta=MagicMock(content="Test"))]),
            MagicMock(choices=[MagicMock(delta=MagicMock(content=" response"))])
        ])
        
        llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
        result = list(llm.stream("test message"))
        
        mock_completion.assert_called_once()
        assert len(result) == 2

    @patch.dict(os.environ, {
        'AWS_ACCESS_KEY_ID': 'test-key-id',
        'AWS_SECRET_ACCESS_KEY': 'test-secret-key',
        'AWS_DEFAULT_REGION': 'us-east-1'
    })
    @patch('litellm.completion')
    def test_bedrock_iam_with_custom_parameters(self, mock_completion):
        """Test Bedrock IAM authentication with custom parameters."""
        mock_completion.return_value = MagicMock()
        mock_completion.return_value.choices = [MagicMock()]
        mock_completion.return_value.choices[0].message.content = "Test response"
        
        llm = LLM(
            model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
            temperature=0.7,
            max_tokens=100
        )
        result = llm.call("test message")
        
        mock_completion.assert_called_once()
        call_args = mock_completion.call_args
        assert call_args[1]['temperature'] == 0.7
        assert call_args[1]['max_tokens'] == 100
        assert result == "Test response"

    @patch.dict(os.environ, {
        'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
        'AWS_DEFAULT_REGION': 'us-west-2'
    })
    @patch('litellm.completion')
    def test_bedrock_api_key_different_region(self, mock_completion):
        """Test Bedrock API key authentication with different region."""
        mock_completion.return_value = MagicMock()
        mock_completion.return_value.choices = [MagicMock()]
        mock_completion.return_value.choices[0].message.content = "Test response"
        
        llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
        result = llm.call("test message")
        
        mock_completion.assert_called_once()
        assert result == "Test response"