import pytest
from awslabs.aws_api_mcp_server.core.aws.driver import (
    IRTranslation,
    get_local_credentials,
    interpret_command,
    translate_cli_to_ir,
)
from awslabs.aws_api_mcp_server.core.common.command import IRCommand
from awslabs.aws_api_mcp_server.core.common.command_metadata import CommandMetadata
from awslabs.aws_api_mcp_server.core.common.errors import (
    DeniedGlobalArgumentsError,
    ExpectedArgumentError,
    InvalidParametersReceivedError,
    InvalidServiceError,
    InvalidServiceOperationError,
    MissingRequiredParametersError,
    ParameterSchemaValidationError,
    ParameterValidationErrorRecord,
    UnknownArgumentsError,
)
from awslabs.aws_api_mcp_server.core.common.models import Credentials
from botocore.exceptions import NoCredentialsError
from tests.fixtures import S3_CLI_NO_REGION, TEST_CREDENTIALS, patch_botocore
from unittest.mock import MagicMock, patch


@patch('awslabs.aws_api_mcp_server.core.aws.driver.boto3.Session')
def test_get_local_credentials_success_with_aws_mcp_profile(mock_session_class):
    """Test get_local_credentials returns credentials when available."""
    mock_session = MagicMock()
    mock_session_class.return_value = mock_session

    mock_credentials = MagicMock()
    mock_credentials.access_key = 'test-access-key'
    mock_credentials.secret_key = 'test-secret-key'  # pragma: allowlist secret
    mock_credentials.token = 'test-session-token'

    mock_session.get_credentials.return_value = mock_credentials

    result = get_local_credentials(profile='test')

    assert isinstance(result, Credentials)
    assert result.access_key_id == 'test-access-key'
    assert result.secret_access_key == 'test-secret-key'  # pragma: allowlist secret
    assert result.session_token == 'test-session-token'
    mock_session_class.assert_called_once_with(profile_name='test')
    mock_session.get_credentials.assert_called_once()


@patch('awslabs.aws_api_mcp_server.core.aws.driver.boto3.Session')
def test_get_local_credentials_success_with_default_creds(mock_session_class):
    """Test get_local_credentials returns credentials when available."""
    mock_session = MagicMock()
    mock_session_class.return_value = mock_session

    mock_credentials = MagicMock()
    mock_credentials.access_key = 'test-access-key'
    mock_credentials.secret_key = 'test-secret-key'  # pragma: allowlist secret
    mock_credentials.token = 'test-session-token'

    mock_session.get_credentials.return_value = mock_credentials

    result = get_local_credentials()

    assert isinstance(result, Credentials)
    assert result.access_key_id == 'test-access-key'
    assert result.secret_access_key == 'test-secret-key'  # pragma: allowlist secret
    assert result.session_token == 'test-session-token'
    mock_session_class.assert_called_once()
    mock_session.get_credentials.assert_called_once()


@patch('awslabs.aws_api_mcp_server.core.aws.driver.boto3.Session')
def test_get_local_credentials_raises_no_credentials_error(mock_session_class):
    """Test get_local_credentials raises NoCredentialsError when credentials are None."""
    mock_session = MagicMock()
    mock_session_class.return_value = mock_session
    mock_session.get_credentials.return_value = None

    with pytest.raises(NoCredentialsError):
        get_local_credentials()

    mock_session_class.assert_called_once()
    mock_session.get_credentials.assert_called_once()


@pytest.mark.parametrize(
    'command,program',
    [
        (
            S3_CLI_NO_REGION,
            IRTranslation(
                command_metadata=CommandMetadata(
                    's3', 'Amazon Simple Storage Service', 'ListBuckets'
                ),
            ),
        ),
        (
            'aws cloud8 list-environments',
            IRTranslation(validation_failures=[InvalidServiceError('cloud8').as_failure()]),
        ),
        # s3 is valid but it is not a real service API - it boils down to multiple API calls to s3api
        (
            'aws s3 ls s3://flock-datasets-us-west-2-516690746032',
            IRTranslation(
                command=IRCommand(
                    command_metadata=CommandMetadata('s3', None, 'ls'),
                    region='us-east-1',
                    parameters={},
                    is_awscli_customization=True,
                ),
                command_metadata=CommandMetadata('s3', None, 'ls'),
            ),
        ),
        (
            'aws s3 ls',
            IRTranslation(
                command=IRCommand(
                    command_metadata=CommandMetadata('s3', None, 'ls'),
                    region='us-east-1',
                    parameters={},
                    is_awscli_customization=True,
                ),
                command_metadata=CommandMetadata('s3', None, 'ls'),
            ),
        ),
        (
            'aws ec2 lss',
            IRTranslation(
                validation_failures=[InvalidServiceOperationError('ec2', 'lss').as_failure()]
            ),
        ),
        (
            'aws cloud9 describe-environment-status',
            IRTranslation(
                missing_context_failures=[
                    MissingRequiredParametersError(
                        'cloud9',
                        'describe-environment-status',
                        ['--environment-id'],
                        CommandMetadata('cloud9', 'AWS Cloud9', 'DescribeEnvironmentStatus'),
                    ).as_failure()
                ],
                command_metadata=CommandMetadata(
                    'cloud9', 'AWS Cloud9', 'DescribeEnvironmentStatus'
                ),
            ),
        ),
        (
            'aws kinesis get-records --shard-iterator',
            IRTranslation(
                missing_context_failures=[
                    ExpectedArgumentError(
                        '--shard-iterator',
                        'expected one argument',
                        CommandMetadata('kinesis', 'Amazon Kinesis', 'GetRecords'),
                    ).as_failure()
                ],
                command_metadata=CommandMetadata('kinesis', 'Amazon Kinesis', 'GetRecords'),
            ),
        ),
        (
            'aws cloud9 describe-environment-status --environment-id xyz --status',
            IRTranslation(
                validation_failures=[
                    InvalidParametersReceivedError(
                        'cloud9',
                        'describe-environment-status',
                        ['--status'],
                        ['--environment-id'],
                    ).as_failure()
                ]
            ),
        ),
        (
            'aws batch list-jobs --no-verify-ssl --debug --no-sign-request',
            IRTranslation(
                validation_failures=[
                    DeniedGlobalArgumentsError(
                        'batch',
                        [
                            '--debug',
                            '--no-sign-request',
                            '--no-verify-ssl',
                        ],
                    ).as_failure()
                ]
            ),
        ),
        (
            'aws s3api get-bucket-intelligent-tiering-configuration --bucket my-bucket --output json',
            IRTranslation(
                missing_context_failures=[
                    MissingRequiredParametersError(
                        's3api',
                        'get-bucket-intelligent-tiering-configuration',
                        ['--id'],
                        CommandMetadata(
                            's3',
                            'Amazon Simple Storage Service',
                            'GetBucketIntelligentTieringConfiguration',
                        ),
                    ).as_failure()
                ],
                command_metadata=CommandMetadata(
                    's3',
                    'Amazon Simple Storage Service',
                    'GetBucketIntelligentTieringConfiguration',
                ),
            ),
        ),
        (
            'aws s3control list-access-grants --account-id test_account',
            IRTranslation(
                validation_failures=[
                    ParameterSchemaValidationError(
                        [
                            ParameterValidationErrorRecord(
                                '--account-id',
                                'Invalid pattern for parameter , value: test_account, valid pattern: ^\\d{12}$',
                            )
                        ]
                    ).as_failure()
                ]
            ),
        ),
        (
            'aws datazone search-listings --domain-identifier dzd_rmvr776t4h0pvi --search-text shipping logistics costs',
            IRTranslation(
                validation_failures=[
                    UnknownArgumentsError(
                        'datazone',
                        'search-listings',
                        ['logistics', 'costs'],
                    ).as_failure()
                ]
            ),
        ),
        (
            'aws kinesis describe-stream --stream-name 12345~**',
            IRTranslation(
                validation_failures=[
                    ParameterSchemaValidationError(
                        [
                            ParameterValidationErrorRecord(
                                '--stream-name',
                                'Invalid pattern for parameter , value: 12345~**, valid pattern: [a-zA-Z0-9_.-]+',
                            )
                        ]
                    ).as_failure()
                ]
            ),
        ),
        # Shape for stream-name has max length 128 but this is not validated to remain forwards compatible with any API changes
        (
            (
                'aws kinesis describe-stream --stream-name 1234511111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111111111111'
                '1111111111111111111111111111111111111111111111111111111111111111111111111'
            ),
            IRTranslation(
                command=IRCommand(
                    command_metadata=CommandMetadata(
                        'kinesis', 'Amazon Kinesis', 'DescribeStream'
                    ),
                    region='us-east-1',
                    parameters={},
                    is_awscli_customization=False,
                ),
                command_metadata=CommandMetadata('kinesis', 'Amazon Kinesis', 'DescribeStream'),
            ),
        ),
    ],
)
def test_driver(command, program):
    """Test that CLI command is correctly translated to IR program."""
    translated = translate_cli_to_ir(command)
    assert translated == program
    assert translated.command_metadata == program.command_metadata


@pytest.mark.parametrize(
    'command',
    [
        # Cloud9 is not available in ap-southeast-3
        'aws cloud9 list-environments --region ap-southeast-3',
        # And fake regions are not tracked as well
        'aws cloud9 list-environments --region bogus',
        # Datazone not available in certain regions
        'aws datazone get-environment --domain-identifier dzd_rm8rqsucr193md --identifier dzd --region eu-central-2',
        'aws datazone get-environment --domain-identifier dzd_rm8rqsucr193md --identifier dzd --region ap-south-1',
        'aws datazone get-environment --domain-identifier dzd_rm8rqsucr193md --identifier dzd --region ap-east-1',
        'aws datazone get-environment --domain-identifier dzd_rm8rqsucr193md --identifier dzd --region eu-central-2',
        # Service is not available in default region
        'aws datazone get-environment --domain-identifier dzd_rm8rqsucr193md --identifier dzd',
    ],
)
def test_invalid_region(command):
    """Test that invalid or unavailable regions are handled correctly."""
    translate_cli_to_ir(command)


# Tests for credentials integration changes
@patch('awslabs.aws_api_mcp_server.core.aws.driver.get_local_credentials')
def test_interpret_command_with_credentials_parameter(mock_get_local_credentials):
    """Test that interpret_command uses provided credentials instead of calling get_local_credentials."""
    # Create test credentials
    test_credentials = Credentials(**TEST_CREDENTIALS)

    # Mock get_local_credentials to ensure it's not called
    mock_get_local_credentials.return_value = test_credentials

    with patch_botocore():
        result = interpret_command('aws s3api list-buckets', credentials=test_credentials)

    # Verify get_local_credentials was not called when credentials were provided
    mock_get_local_credentials.assert_not_called()
    assert result is not None


@patch('awslabs.aws_api_mcp_server.core.aws.driver.get_local_credentials')
def test_interpret_command_without_credentials_parameter(mock_get_local_credentials):
    """Test that interpret_command falls back to get_local_credentials when no credentials provided."""
    # Create test credentials
    test_credentials = Credentials(**TEST_CREDENTIALS)

    mock_get_local_credentials.return_value = test_credentials

    with patch_botocore():
        result = interpret_command('aws s3api list-buckets')

    # Verify get_local_credentials was called when no credentials were provided
    mock_get_local_credentials.assert_called_once()
    assert result is not None


@patch('awslabs.aws_api_mcp_server.core.aws.driver.get_local_credentials')
def test_interpret_command_credentials_precedence(mock_get_local_credentials):
    """Test that provided credentials take precedence over local credentials."""
    # Create different credentials to test precedence
    local_credentials = Credentials(**TEST_CREDENTIALS)

    provided_credentials = Credentials(**TEST_CREDENTIALS)

    mock_get_local_credentials.return_value = local_credentials

    with patch_botocore():
        result = interpret_command('aws s3api list-buckets', credentials=provided_credentials)

    # Verify get_local_credentials was not called when credentials were provided
    mock_get_local_credentials.assert_not_called()
    assert result is not None
