"""Tests for the CLI executor module."""

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from aws_mcp_server.cli_executor import (
    CommandExecutionError,
    CommandValidationError,
    check_aws_cli_installed,
    execute_aws_command,
    execute_pipe_command,
    get_command_help,
    is_auth_error,
)
from aws_mcp_server.config import DEFAULT_TIMEOUT, MAX_OUTPUT_SIZE


@pytest.mark.asyncio
async def test_execute_aws_command_success():
    """Test successful command execution."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a successful process
        process_mock = AsyncMock()
        process_mock.returncode = 0
        process_mock.communicate.return_value = (b"Success output", b"")
        mock_subprocess.return_value = process_mock

        result = await execute_aws_command("aws s3 ls")

        assert result["status"] == "success"
        assert result["output"] == "Success output"
        mock_subprocess.assert_called_once_with("aws", "s3", "ls", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)


@pytest.mark.asyncio
async def test_execute_aws_command_ec2_with_region_added():
    """Test that region is automatically added to EC2 commands."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a successful process
        process_mock = AsyncMock()
        process_mock.returncode = 0
        process_mock.communicate.return_value = (b"EC2 instances", b"")
        mock_subprocess.return_value = process_mock

        # Import here to ensure the test uses the actual value
        from aws_mcp_server.config import AWS_REGION

        # Execute an EC2 command without region
        result = await execute_aws_command("aws ec2 describe-instances")

        assert result["status"] == "success"
        assert result["output"] == "EC2 instances"

        # Verify region was added to the command
        mock_subprocess.assert_called_once()
        call_args = mock_subprocess.call_args[0]
        assert call_args[0] == "aws"
        assert call_args[1] == "ec2"
        assert call_args[2] == "describe-instances"
        assert "--region" in call_args
        assert AWS_REGION in call_args


@pytest.mark.asyncio
async def test_execute_aws_command_with_custom_timeout():
    """Test command execution with custom timeout."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        process_mock = AsyncMock()
        process_mock.returncode = 0
        process_mock.communicate.return_value = (b"Success output", b"")
        mock_subprocess.return_value = process_mock

        # Use a custom timeout
        custom_timeout = 120
        with patch("asyncio.wait_for") as mock_wait_for:
            mock_wait_for.return_value = (b"Success output", b"")
            await execute_aws_command("aws s3 ls", timeout=custom_timeout)

            # Check that wait_for was called with the custom timeout
            mock_wait_for.assert_called_once()
            args, kwargs = mock_wait_for.call_args
            assert kwargs.get("timeout") == custom_timeout or args[1] == custom_timeout


@pytest.mark.asyncio
async def test_execute_aws_command_error():
    """Test command execution error."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a failed process
        process_mock = AsyncMock()
        process_mock.returncode = 1
        # Set up an awaitable communicate method
        communicate_mock = AsyncMock()
        communicate_mock.return_value = (b"", b"Error message")
        process_mock.communicate = communicate_mock
        mock_subprocess.return_value = process_mock

        result = await execute_aws_command("aws s3 ls")

        assert result["status"] == "error"
        assert result["output"] == "Error message"
        # Verify communicate was called
        communicate_mock.assert_called_once()


@pytest.mark.asyncio
async def test_execute_aws_command_auth_error():
    """Test command execution with authentication error."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a process that returns auth error
        process_mock = AsyncMock()
        process_mock.returncode = 1
        process_mock.communicate.return_value = (b"", b"Unable to locate credentials")
        mock_subprocess.return_value = process_mock

        result = await execute_aws_command("aws s3 ls")

        assert result["status"] == "error"
        assert "Authentication error" in result["output"]
        assert "Unable to locate credentials" in result["output"]
        assert "Please check your AWS credentials" in result["output"]


@pytest.mark.asyncio
async def test_execute_aws_command_timeout():
    """Test command timeout."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a process that times out
        process_mock = AsyncMock()
        # Use a properly awaitable mock that raises TimeoutError
        communicate_mock = AsyncMock(side_effect=asyncio.TimeoutError())
        process_mock.communicate = communicate_mock
        mock_subprocess.return_value = process_mock

        # Mock a regular function instead of an async one for process.kill
        process_mock.kill = MagicMock()

        with pytest.raises(CommandExecutionError) as excinfo:
            await execute_aws_command("aws s3 ls", timeout=1)

        # Check error message
        assert "Command timed out after 1 seconds" in str(excinfo.value)

        # Verify process was killed
        process_mock.kill.assert_called_once()


@pytest.mark.asyncio
async def test_execute_aws_command_kill_failure():
    """Test failure to kill process after timeout."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a process that times out
        process_mock = AsyncMock()
        # Use a properly awaitable mock that raises TimeoutError
        communicate_mock = AsyncMock(side_effect=asyncio.TimeoutError())
        process_mock.communicate = communicate_mock
        # Use regular MagicMock since kill() is not an async method
        process_mock.kill = MagicMock(side_effect=Exception("Failed to kill process"))
        mock_subprocess.return_value = process_mock

        with pytest.raises(CommandExecutionError) as excinfo:
            await execute_aws_command("aws s3 ls", timeout=1)

        # The main exception should still be about the timeout
        assert "Command timed out after 1 seconds" in str(excinfo.value)


@pytest.mark.asyncio
async def test_execute_aws_command_general_exception():
    """Test handling of general exceptions during command execution."""
    with patch("asyncio.create_subprocess_exec", side_effect=Exception("Test exception")):
        with pytest.raises(CommandExecutionError) as excinfo:
            await execute_aws_command("aws s3 ls")

        assert "Failed to execute command" in str(excinfo.value)
        assert "Test exception" in str(excinfo.value)


@pytest.mark.asyncio
async def test_execute_aws_command_truncate_output():
    """Test truncation of large outputs."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        # Mock a successful process with large output
        process_mock = AsyncMock()
        process_mock.returncode = 0

        # Generate a large output that exceeds MAX_OUTPUT_SIZE
        large_output = "x" * (MAX_OUTPUT_SIZE + 1000)
        process_mock.communicate.return_value = (large_output.encode("utf-8"), b"")
        mock_subprocess.return_value = process_mock

        result = await execute_aws_command("aws s3 ls")

        assert result["status"] == "success"
        assert len(result["output"]) <= MAX_OUTPUT_SIZE + 100  # Allow for the truncation message
        assert "output truncated" in result["output"]


@pytest.mark.parametrize(
    "error_message,expected_result",
    [
        # Positive cases
        ("Unable to locate credentials", True),
        ("Some text before ExpiredToken and after", True),
        ("Error: AccessDenied when attempting to perform operation", True),
        ("AuthFailure: credentials could not be verified", True),
        ("The security token included in the request is invalid", True),
        ("The config profile could not be found", True),
        # Negative cases
        ("S3 bucket not found", False),
        ("Resource not found: myresource", False),
        ("Invalid parameter value", False),
    ],
)
def test_is_auth_error(error_message, expected_result):
    """Test the is_auth_error function with various error messages."""
    assert is_auth_error(error_message) == expected_result


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "returncode,stdout,stderr,exception,expected_result",
    [
        # CLI installed
        (0, b"aws-cli/2.15.0", b"", None, True),
        # CLI not installed - command not found
        (127, b"", b"command not found", None, False),
        # CLI error case
        (1, b"", b"some error", None, False),
        # Exception during command execution
        (None, None, None, Exception("Test exception"), False),
    ],
)
async def test_check_aws_cli_installed(returncode, stdout, stderr, exception, expected_result):
    """Test check_aws_cli_installed function with various scenarios."""
    if exception:
        with patch("asyncio.create_subprocess_exec", side_effect=exception):
            result = await check_aws_cli_installed()
            assert result is expected_result
    else:
        with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
            process_mock = AsyncMock()
            process_mock.returncode = returncode
            process_mock.communicate.return_value = (stdout, stderr)
            mock_subprocess.return_value = process_mock

            result = await check_aws_cli_installed()
            assert result is expected_result

            if returncode == 0:  # Only verify call args for success case to avoid redundancy
                mock_subprocess.assert_called_once_with("aws", "--version", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "service,command,mock_type,mock_value,expected_text,expected_call",
    [
        # Successful help retrieval with service and command
        ("s3", "ls", "return_value", {"status": "success", "output": "Help text"}, "Help text", "aws s3 ls help"),
        # Successful help retrieval with service only
        ("s3", None, "return_value", {"status": "success", "output": "Help text for service"}, "Help text for service", "aws s3 help"),
        # Error scenarios
        ("s3", "ls", "side_effect", CommandValidationError("Test validation error"), "Command validation error: Test validation error", None),
        ("s3", "ls", "side_effect", CommandExecutionError("Test execution error"), "Error retrieving help: Test execution error", None),
        ("s3", "ls", "side_effect", Exception("Test exception"), "Error retrieving help: Test exception", None),
        # Error result from AWS command
        ("s3", "ls", "return_value", {"status": "error", "output": "Command failed"}, "Error: Command failed", "aws s3 ls help"),
    ],
)
async def test_get_command_help(service, command, mock_type, mock_value, expected_text, expected_call):
    """Test get_command_help function with various scenarios."""
    with patch("aws_mcp_server.cli_executor.execute_aws_command", new_callable=AsyncMock) as mock_execute:
        # Configure the mock based on the test case
        if mock_type == "return_value":
            mock_execute.return_value = mock_value
        else:  # side_effect
            mock_execute.side_effect = mock_value

        # Call the function
        result = await get_command_help(service, command)

        # Verify the result
        assert expected_text in result["help_text"]

        # Verify the mock was called correctly if expected_call is provided
        if expected_call:
            mock_execute.assert_called_once_with(expected_call)


@pytest.mark.asyncio
async def test_execute_aws_command_with_pipe():
    """Test execute_aws_command with a piped command."""
    # Test that execute_aws_command calls execute_pipe_command for piped commands
    with patch("aws_mcp_server.cli_executor.is_pipe_command", return_value=True):
        with patch("aws_mcp_server.cli_executor.execute_pipe_command", new_callable=AsyncMock) as mock_pipe_exec:
            mock_pipe_exec.return_value = {"status": "success", "output": "Piped result"}

            result = await execute_aws_command("aws s3 ls | grep bucket")

            assert result["status"] == "success"
            assert result["output"] == "Piped result"
            mock_pipe_exec.assert_called_once_with("aws s3 ls | grep bucket", None)


@pytest.mark.asyncio
async def test_execute_pipe_command_success():
    """Test successful execution of a pipe command."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command") as mock_validate:
        with patch("aws_mcp_server.cli_executor.execute_piped_command", new_callable=AsyncMock) as mock_pipe_exec:
            mock_pipe_exec.return_value = {"status": "success", "output": "Filtered results"}

            result = await execute_pipe_command("aws s3 ls | grep bucket")

            assert result["status"] == "success"
            assert result["output"] == "Filtered results"
            mock_validate.assert_called_once_with("aws s3 ls | grep bucket")
            mock_pipe_exec.assert_called_once_with("aws s3 ls | grep bucket", None)


@pytest.mark.asyncio
async def test_execute_pipe_command_ec2_with_region_added():
    """Test that region is automatically added to EC2 commands in a pipe."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command"):
        with patch("aws_mcp_server.cli_executor.execute_piped_command", new_callable=AsyncMock) as mock_pipe_exec:
            mock_pipe_exec.return_value = {"status": "success", "output": "Filtered EC2 instances"}

            # Mock split_pipe_command to simulate pipe command splitting
            with patch("aws_mcp_server.cli_executor.split_pipe_command") as mock_split:
                mock_split.return_value = ["aws ec2 describe-instances", "grep instance-id"]

                # Import here to ensure the test uses the actual value
                from aws_mcp_server.config import AWS_REGION

                # Execute a piped EC2 command without region
                result = await execute_pipe_command("aws ec2 describe-instances | grep instance-id")

                assert result["status"] == "success"
                assert result["output"] == "Filtered EC2 instances"

                # Verify the command was modified to include region
                expected_cmd = f"aws ec2 describe-instances --region {AWS_REGION} | grep instance-id"
                mock_pipe_exec.assert_called_once_with(expected_cmd, None)


@pytest.mark.asyncio
async def test_execute_pipe_command_validation_error():
    """Test execute_pipe_command with validation error."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command", side_effect=CommandValidationError("Invalid pipe command")):
        with pytest.raises(CommandValidationError) as excinfo:
            await execute_pipe_command("invalid | pipe | command")

        assert "Invalid pipe command" in str(excinfo.value)


@pytest.mark.asyncio
async def test_execute_pipe_command_execution_error():
    """Test execute_pipe_command with execution error."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command"):
        with patch("aws_mcp_server.cli_executor.execute_piped_command", side_effect=Exception("Execution error")):
            with pytest.raises(CommandExecutionError) as excinfo:
                await execute_pipe_command("aws s3 ls | grep bucket")

            assert "Failed to execute piped command" in str(excinfo.value)
            assert "Execution error" in str(excinfo.value)


# New test cases to improve coverage


@pytest.mark.asyncio
async def test_execute_pipe_command_timeout():
    """Test timeout handling in piped commands."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command"):
        with patch("aws_mcp_server.cli_executor.execute_piped_command", new_callable=AsyncMock) as mock_exec:
            # Simulate timeout in the executed command
            mock_exec.return_value = {"status": "error", "output": f"Command timed out after {DEFAULT_TIMEOUT} seconds"}

            result = await execute_pipe_command("aws s3 ls | grep bucket")

            assert result["status"] == "error"
            assert f"Command timed out after {DEFAULT_TIMEOUT} seconds" in result["output"]
            mock_exec.assert_called_once()


@pytest.mark.asyncio
async def test_execute_pipe_command_with_custom_timeout():
    """Test piped command execution with custom timeout."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command"):
        with patch("aws_mcp_server.cli_executor.execute_piped_command", new_callable=AsyncMock) as mock_exec:
            mock_exec.return_value = {"status": "success", "output": "Piped output"}

            custom_timeout = 120
            await execute_pipe_command("aws s3 ls | grep bucket", timeout=custom_timeout)

            # Verify the custom timeout was passed to the execute_piped_command
            mock_exec.assert_called_once_with("aws s3 ls | grep bucket", custom_timeout)


@pytest.mark.asyncio
async def test_execute_pipe_command_large_output():
    """Test handling of large output in piped commands."""
    with patch("aws_mcp_server.cli_executor.validate_pipe_command"):
        with patch("aws_mcp_server.cli_executor.execute_piped_command", new_callable=AsyncMock) as mock_exec:
            # Generate large output that would be truncated
            large_output = "x" * (MAX_OUTPUT_SIZE + 1000)
            mock_exec.return_value = {"status": "success", "output": large_output}

            result = await execute_pipe_command("aws s3 ls | grep bucket")

            assert result["status"] == "success"
            assert len(result["output"]) == len(large_output)  # Length should be preserved here as truncation happens in tools module


@pytest.mark.parametrize(
    "exit_code,stderr,expected_status,expected_msg",
    [
        (0, b"", "success", ""),  # Success case
        (1, b"Error: bucket not found", "error", "Error: bucket not found"),  # Standard error
        (1, b"AccessDenied", "error", "Authentication error"),  # Auth error
        (0, b"Warning: deprecated feature", "success", ""),  # Warning on stderr but success exit code
    ],
)
@pytest.mark.asyncio
async def test_execute_aws_command_exit_codes(exit_code, stderr, expected_status, expected_msg):
    """Test handling of different process exit codes and stderr output."""
    with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess:
        process_mock = AsyncMock()
        process_mock.returncode = exit_code
        stdout = b"Command output" if exit_code == 0 else b""
        process_mock.communicate.return_value = (stdout, stderr)
        mock_subprocess.return_value = process_mock

        result = await execute_aws_command("aws s3 ls")

        assert result["status"] == expected_status
        if expected_status == "success":
            assert result["output"] == "Command output"
        else:
            assert expected_msg in result["output"]
