"""
Comprehensive test suite for the UnloadDataSource class.

This test class validates the Redshift UNLOAD data source implementation,
including connection handling, statement execution, database switching,
and credential sanitization.
"""

import unittest
from unittest.mock import MagicMock, Mock, patch

from data_exchange_agent.data_sources.sql_command_type import SQLCommandType
from data_exchange_agent.data_sources.unload_data_source import (
    UnloadDataSource,
    _sanitize_unload_statement,
)


class TestSanitizeUnloadStatement(unittest.TestCase):
    """
    Test suite for the _sanitize_unload_statement function.

    Validates that IAM role ARNs are properly redacted from UNLOAD statements
    to prevent credential leakage in logs.
    """

    def test_sanitize_single_quoted_iam_role(self):
        """Test sanitization of single-quoted IAM role ARN."""
        statement = "UNLOAD ('SELECT * FROM users') " "TO 's3://bucket/path/' " "IAM_ROLE 'arn:aws:iam::123456789012:role/MyRedshiftRole'"
        result = _sanitize_unload_statement(statement)
        self.assertIn("IAM_ROLE '***'", result)
        self.assertNotIn("123456789012", result)
        self.assertNotIn("MyRedshiftRole", result)

    def test_sanitize_double_quoted_iam_role(self):
        """Test sanitization of double-quoted IAM role ARN."""
        statement = "UNLOAD ('SELECT * FROM users') " "TO 's3://bucket/path/' " 'IAM_ROLE "arn:aws:iam::123456789012:role/MyRedshiftRole"'
        result = _sanitize_unload_statement(statement)
        self.assertIn("IAM_ROLE '***'", result)
        self.assertNotIn("123456789012", result)

    def test_sanitize_preserves_other_content(self):
        """Test that sanitization preserves non-sensitive content."""
        statement = (
            "UNLOAD ('SELECT id, name FROM schema.table WHERE active = true') "
            "TO 's3://my-bucket/exports/data/' "
            "IAM_ROLE 'arn:aws:iam::999:role/Role' "
            "FORMAT AS PARQUET MAXFILESIZE 256 MB"
        )
        result = _sanitize_unload_statement(statement)

        # Verify sensitive data is redacted
        self.assertIn("IAM_ROLE '***'", result)

        # Verify non-sensitive data is preserved
        self.assertIn("SELECT id, name FROM schema.table WHERE active = true", result)
        self.assertIn("s3://my-bucket/exports/data/", result)
        self.assertIn("FORMAT AS PARQUET", result)
        self.assertIn("MAXFILESIZE 256 MB", result)

    def test_sanitize_case_insensitive_iam_role(self):
        """Test that IAM_ROLE matching is case insensitive."""
        statements = [
            "UNLOAD ('SELECT 1') TO 's3://b/' iam_role 'arn:aws:iam::1:role/R'",
            "UNLOAD ('SELECT 1') TO 's3://b/' Iam_Role 'arn:aws:iam::1:role/R'",
            "UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
        ]

        for statement in statements:
            result = _sanitize_unload_statement(statement)
            self.assertNotIn("arn:aws:iam::", result, f"Failed for: {statement}")

    def test_sanitize_statement_without_iam_role(self):
        """Test that statements without IAM_ROLE are unchanged."""
        statement = "SELECT * FROM users WHERE id = 123"
        result = _sanitize_unload_statement(statement)
        self.assertEqual(result, statement)

    def test_sanitize_complex_arn(self):
        """Test sanitization of complex IAM role ARN with path."""
        statement = (
            "UNLOAD ('SELECT * FROM t') TO 's3://b/' " "IAM_ROLE 'arn:aws:iam::123456789012:role/service-role/RedshiftCopyUnload-role'"
        )
        result = _sanitize_unload_statement(statement)
        self.assertIn("IAM_ROLE '***'", result)
        self.assertNotIn("RedshiftCopyUnload", result)


class TestUnloadDataSourceProperties(unittest.TestCase):
    """
    Test suite for UnloadDataSource properties.

    Validates that the data source properties are correctly initialized
    and accessible.
    """

    @patch("data_exchange_agent.data_sources.unload_data_source.SFLogger")
    def test_statement_property(self, mock_logger_class):
        """Test that statement property returns the UNLOAD statement."""
        mock_logger = Mock()
        mock_program_config = Mock()

        # Setup mock connection config
        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        unload_statement = "UNLOAD ('SELECT * FROM users') " "TO 's3://bucket/path/' " "IAM_ROLE 'arn:aws:iam::123:role/Role'"

        data_source = UnloadDataSource(
            engine="redshift",
            statement=unload_statement,
            results_folder_path="/tmp/results",
            base_file_name="result",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        self.assertEqual(data_source.statement, unload_statement)

    @patch("data_exchange_agent.data_sources.unload_data_source.SFLogger")
    def test_results_folder_path_property(self, mock_logger_class):
        """Test that results_folder_path property is set correctly."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            results_folder_path="/custom/path",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        self.assertEqual(data_source.results_folder_path, "/custom/path")

    @patch("data_exchange_agent.data_sources.unload_data_source.SFLogger")
    def test_results_folder_path_defaults_to_empty(self, mock_logger_class):
        """Test that results_folder_path defaults to empty string when None."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            results_folder_path=None,
            logger=mock_logger,
            program_config=mock_program_config,
        )

        self.assertEqual(data_source.results_folder_path, "")

    @patch("data_exchange_agent.data_sources.unload_data_source.SFLogger")
    def test_base_file_name_property(self, mock_logger_class):
        """Test that base_file_name property is set correctly."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            base_file_name="custom_file",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        self.assertEqual(data_source.base_file_name, "custom_file")


class TestUnloadDataSourceExportData(unittest.TestCase):
    """
    Test suite for UnloadDataSource.export_data method.

    Validates the UNLOAD command execution, including connection handling,
    statement validation, and error handling.
    """

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_executes_unload_command(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that export_data executes the UNLOAD command successfully."""
        # Setup mocks
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        unload_statement = "UNLOAD ('SELECT * FROM users') " "TO 's3://bucket/path/' " "IAM_ROLE 'arn:aws:iam::123:role/Role'"

        data_source = UnloadDataSource(
            engine="redshift",
            statement=unload_statement,
            logger=mock_logger,
            program_config=mock_program_config,
        )

        result = data_source.export_data()

        self.assertTrue(result)
        mock_pyodbc.connect.assert_called_once()
        mock_cursor.execute.assert_called_once_with(unload_statement)
        mock_cursor.close.assert_called_once()
        mock_connection.close.assert_called_once()

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_rejects_non_unload_statement(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that export_data raises exception for non-UNLOAD statements."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        # Return SELECT type instead of UNLOAD
        mock_get_sql_type.return_value = SQLCommandType.SELECT

        data_source = UnloadDataSource(
            engine="redshift",
            statement="SELECT * FROM users",  # Not an UNLOAD statement
            logger=mock_logger,
            program_config=mock_program_config,
        )

        with self.assertRaises(Exception) as context:
            data_source.export_data()

        self.assertIn("UnloadDataSource can only execute UNLOAD commands", str(context.exception))

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_handles_connection_error(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that export_data handles connection errors."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_pyodbc.connect.side_effect = Exception("Connection failed")
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        with self.assertRaises(Exception) as context:
            data_source.export_data()

        self.assertIn("Connection failed", str(context.exception))
        mock_logger.error.assert_called()

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_handles_execution_error(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that export_data handles statement execution errors."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_cursor.execute.side_effect = Exception("Permission denied")
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        with self.assertRaises(Exception) as context:
            data_source.export_data()

        self.assertIn("Permission denied", str(context.exception))
        mock_logger.error.assert_called()
        # Connection should be closed even on error
        mock_connection.close.assert_called_once()

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_closes_connection_on_success(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that connection is closed after successful execution."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        data_source.export_data()

        mock_connection.close.assert_called_once()


class TestUnloadDataSourceDatabaseSwitching(unittest.TestCase):
    """
    Test suite for UnloadDataSource database switching functionality.

    Validates the _switch_database method and use_database_command handling.
    """

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_switches_database_when_command_provided(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that database switch is executed when use_database_command is provided."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        unload_statement = "UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'"

        data_source = UnloadDataSource(
            engine="redshift",
            statement=unload_statement,
            use_database_command="USE my_database",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        data_source.export_data()

        # Verify cursor was used for both USE command and UNLOAD
        calls = mock_cursor.execute.call_args_list
        self.assertEqual(len(calls), 2)
        self.assertEqual(calls[0][0][0], "USE my_database")
        self.assertEqual(calls[1][0][0], unload_statement)

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_export_data_skips_database_switch_when_no_command(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that no database switch occurs when use_database_command is None."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        unload_statement = "UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'"

        data_source = UnloadDataSource(
            engine="redshift",
            statement=unload_statement,
            use_database_command=None,
            logger=mock_logger,
            program_config=mock_program_config,
        )

        data_source.export_data()

        # Verify cursor was only used for UNLOAD
        calls = mock_cursor.execute.call_args_list
        self.assertEqual(len(calls), 1)
        self.assertEqual(calls[0][0][0], unload_statement)

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_database_switch_error_propagates(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that database switch errors are properly propagated."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        # Make the first execute (USE command) fail
        mock_cursor.execute.side_effect = Exception("Database not found")
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        data_source = UnloadDataSource(
            engine="redshift",
            statement="UNLOAD ('SELECT 1') TO 's3://b/' IAM_ROLE 'arn:aws:iam::1:role/R'",
            use_database_command="USE nonexistent_db",
            logger=mock_logger,
            program_config=mock_program_config,
        )

        with self.assertRaises(Exception) as context:
            data_source.export_data()

        self.assertIn("Database not found", str(context.exception))
        mock_logger.error.assert_called()


class TestUnloadDataSourceLogging(unittest.TestCase):
    """
    Test suite for UnloadDataSource logging behavior.

    Validates that sensitive information is properly redacted in logs.
    """

    @patch("data_exchange_agent.data_sources.unload_data_source._get_pyodbc")
    @patch("data_exchange_agent.data_sources.unload_data_source.get_permitted_sql_command_type")
    def test_logs_sanitized_statement(self, mock_get_sql_type, mock_get_pyodbc):
        """Test that logged statements have IAM roles redacted."""
        mock_logger = Mock()
        mock_program_config = Mock()

        mock_connection_config = Mock()
        mock_connection_config.connection_string = "Driver={Test};Server=test"
        mock_program_config.__getitem__ = Mock(return_value={"redshift": mock_connection_config})

        mock_pyodbc = MagicMock()
        mock_connection = MagicMock()
        mock_cursor = MagicMock()
        mock_pyodbc.connect.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        mock_get_pyodbc.return_value = mock_pyodbc
        mock_get_sql_type.return_value = SQLCommandType.UNLOAD

        iam_role_arn = "arn:aws:iam::123456789012:role/SensitiveRole"
        unload_statement = f"UNLOAD ('SELECT * FROM users') TO 's3://bucket/' IAM_ROLE '{iam_role_arn}'"

        data_source = UnloadDataSource(
            engine="redshift",
            statement=unload_statement,
            logger=mock_logger,
            program_config=mock_program_config,
        )

        data_source.export_data()

        # Check that logger.info was called and IAM role was NOT in any of the log messages
        for call in mock_logger.info.call_args_list:
            log_message = str(call)
            self.assertNotIn(iam_role_arn, log_message)
            self.assertNotIn("123456789012", log_message)
            self.assertNotIn("SensitiveRole", log_message)


class TestUnloadDataSourceRegistration(unittest.TestCase):
    """
    Test suite for UnloadDataSource registration in the data source registry.

    Validates that UnloadDataSource is properly registered and can be created
    through the registry.
    """

    def test_unload_data_source_is_registered(self):
        """Test that UnloadDataSource is registered in the registry."""
        from data_exchange_agent.constants.data_source_types import DataSourceType
        from data_exchange_agent.data_sources import DataSourceRegistry

        self.assertTrue(DataSourceRegistry.is_registered(DataSourceType.UNLOAD))

    def test_get_unload_data_source_class(self):
        """Test getting the UnloadDataSource class from registry."""
        from data_exchange_agent.constants.data_source_types import DataSourceType
        from data_exchange_agent.data_sources import DataSourceRegistry, UnloadDataSource

        result = DataSourceRegistry.get(DataSourceType.UNLOAD)

        self.assertEqual(result, UnloadDataSource)


if __name__ == "__main__":
    unittest.main()
