"""Unit tests for the Airflow MWAA connector."""

import base64
from unittest.mock import MagicMock, patch

import pytest

from abacus_contract.connectors import airflow
from core.config import Config


def test_sanitize_error_message_filters_warnings_and_keeps_errors():
    """Ensure warnings are stripped, errors remain."""
    raw = (
        'Warning: transient issue\n'
        'some warning happened\n'
        'getattr(module, class_name) failed\n'
        'ERROR: Boom\n'
        '\n'
    )
    encoded = base64.b64encode(raw.encode('utf-8')).decode('utf-8')

    out = airflow._sanitize_error_message(encoded)
    assert out == 'ERROR: Boom'


def test_sanitize_error_message_all_filtered_becomes_empty_string():
    """All-warning input becomes empty string output."""
    raw = 'WARNING: ignore me\ngetattr(module, class_name) blah\n\n'
    encoded = base64.b64encode(raw.encode('utf-8')).decode('utf-8')

    out = airflow._sanitize_error_message(encoded)
    assert out == ''


def test_sanitize_error_message_handles_empty_base64():
    """Empty base64 input becomes empty string output."""
    out = airflow._sanitize_error_message('')  # base64 of empty -> ""
    assert out == ''


def test_run_airflow_cli_command_success_no_stderr():
    """Simulate successful CLI command with no stderr."""
    with (
        patch('abacus_contract.connectors.airflow.boto3.client') as mock_boto_client,
        patch('abacus_contract.connectors.airflow.requests.post') as mock_post,
    ):
        mock_mwaa = MagicMock()
        mock_boto_client.return_value = mock_mwaa
        mock_mwaa.create_cli_token.return_value = {
            'CliToken': 'cli-token-123',
            'WebServerHostname': 'mwaa.example.com',
        }

        # HTTP response: stdout only
        stdout_b64 = base64.b64encode(b'CLI OK').decode('utf-8')
        stderr_b64 = ''  # base64 of empty
        mock_resp = MagicMock()
        mock_resp.json.return_value = {'stdout': stdout_b64, 'stderr': stderr_b64}
        mock_post.return_value = mock_resp

        cmd = 'dags list-runs -d example -s 2025-09-01T00:00:00Z -e 2025-09-26T00:00:00Z --state success --output json'
        result = airflow.run_airflow_cli_command(cmd)

        # Token + URL + headers + body
        mock_boto_client.assert_called_once_with('mwaa')
        mock_mwaa.create_cli_token.assert_called_once_with(Name='qa-abacus-airflow')

        mock_post.assert_called_once()
        url_arg = mock_post.call_args[0][0]
        kwargs = mock_post.call_args[1]
        assert url_arg == 'https://mwaa.example.com/aws_mwaa/cli'
        assert kwargs['data'] == cmd
        assert kwargs['headers']['Authorization'] == 'Bearer cli-token-123'
        assert kwargs['headers']['Content-Type'] == 'text/plain'

        assert result == {'output': 'CLI OK', 'errors': None}


def test_run_airflow_cli_command_filters_warning_stderr_to_none():
    """Simulate stderr with only warnings; ensure we get None errors back."""
    with (
        patch('abacus_contract.connectors.airflow.boto3.client') as mock_boto_client,
        patch('abacus_contract.connectors.airflow.requests.post') as mock_post,
    ):
        mock_mwaa = MagicMock()
        mock_boto_client.return_value = mock_mwaa
        mock_mwaa.create_cli_token.return_value = {
            'CliToken': 'cli-token-xyz',
            'WebServerHostname': 'host.foo',
        }

        stdout_b64 = base64.b64encode(b'OK').decode('utf-8')
        stderr_b64 = base64.b64encode(b'Warning: noisy but harmless').decode('utf-8')
        mock_resp = MagicMock()
        mock_resp.json.return_value = {'stdout': stdout_b64, 'stderr': stderr_b64}
        mock_post.return_value = mock_resp

        result = airflow.run_airflow_cli_command('echo test')

        assert result['output'] == 'OK'
        assert result['errors'] is None  # warnings removed -> treated as no error


def test_run_airflow_cli_command_returns_sanitized_errors_when_present():
    """Simulate stderr with warnings + real error; ensure warnings are stripped."""
    with (
        patch('abacus_contract.connectors.airflow.boto3.client') as mock_boto_client,
        patch('abacus_contract.connectors.airflow.requests.post') as mock_post,
    ):
        mock_mwaa = MagicMock()
        mock_boto_client.return_value = mock_mwaa
        mock_mwaa.create_cli_token.return_value = {
            'CliToken': 't',
            'WebServerHostname': 'h',
        }

        stdout_b64 = base64.b64encode(b'partial output').decode('utf-8')
        stderr_raw = 'warning: flaky thing\nERROR: task failed\n'
        stderr_b64 = base64.b64encode(stderr_raw.encode('utf-8')).decode('utf-8')

        mock_resp = MagicMock()
        mock_resp.json.return_value = {'stdout': stdout_b64, 'stderr': stderr_b64}
        mock_post.return_value = mock_resp

        result = airflow.run_airflow_cli_command('airflow tasks run ...')

        assert result['output'] == 'partial output'
        # warning line filtered; error line remains
        assert result['errors'] == 'ERROR: task failed'


def test_run_airflow_cli_command_uses_env_override_for_mwaa_name():
    """Ensure AIRFLOW_ENV_NAME is respected at *import time*. We re-load the module with the env var set, then run a trivial call."""
    # Set env var before reloading
    with patch.object(Config, 'ENVIRONMENT', Config.PROD_ENVIRONMENT):
        # Reload module to re-evaluate airflow_env_name)
        with (
            patch(
                'abacus_contract.connectors.airflow.boto3.client'
            ) as mock_boto_client,
            patch('abacus_contract.connectors.airflow.requests.post') as mock_post,
        ):
            mock_mwaa = MagicMock()
            mock_boto_client.return_value = mock_mwaa
            mock_mwaa.create_cli_token.return_value = {
                'CliToken': 'abc',
                'WebServerHostname': 'env.host',
            }

            stdout_b64 = base64.b64encode(b'hi').decode('utf-8')
            mock_resp = MagicMock()
            mock_resp.json.return_value = {'stdout': stdout_b64, 'stderr': ''}
            mock_post.return_value = mock_resp

            _ = airflow.run_airflow_cli_command('echo hi')

            mock_mwaa.create_cli_token.assert_called_once_with(
                Name='prod-abacus-airflow'
            )