"""
Tests for TaskManager UNLOAD-related functionality.

This test class validates the UNLOAD statement preparation and data source
resolution logic in the TaskManager.

Note: These tests require the full data-exchange-agent environment to be
properly configured with all dependencies including the shared module.
"""

import pytest
from unittest.mock import Mock, patch

# Check if required modules are available
try:
    from data_exchange_agent.tasks.manager import TaskManager

    TASK_MANAGER_AVAILABLE = True
except ImportError:
    TASK_MANAGER_AVAILABLE = False

try:
    from data_exchange_agent.config.sections.connections.odbc.sqlserver import (
        SQLServerConnectionConfig,
    )

    SQLSERVER_CONFIG_AVAILABLE = True
except ImportError:
    SQLSERVER_CONFIG_AVAILABLE = False

try:
    from data_exchange_agent.config.sections.connections.odbc.redshift import (
        RedshiftConnectionConfig,
        StandardCredentials,
        UnloadConfig,
    )

    REDSHIFT_CONFIG_AVAILABLE = True
except ImportError:
    REDSHIFT_CONFIG_AVAILABLE = False


def _create_task_manager(source_connections_config=None, mock_logger=None):
    """Create a TaskManager with mocked dependencies."""
    mock_logger = mock_logger or Mock()

    if source_connections_config is None:
        source_connections_config = {}

    with patch.object(TaskManager, "__init__", lambda self, **kwargs: None):
        manager = TaskManager()
        manager.logger = mock_logger
        manager.source_connections_config = source_connections_config
        manager.target_connections_config = {}
        manager.task_source_adapter = Mock()
        manager.snowflake_datasource = Mock()
        manager.lease_refresher = Mock()

    return manager


@pytest.mark.skipif(
    not TASK_MANAGER_AVAILABLE or not REDSHIFT_CONFIG_AVAILABLE,
    reason="TaskManager or Redshift config not available (missing shared module)",
)
class TestTaskManagerPrepareUnloadStatement:
    """
    Test suite for TaskManager._prepare_unload_statement method.

    Validates the construction of UNLOAD statements from SELECT queries
    and DEA configuration.
    """

    def test_prepare_unload_statement_success(self):
        """Test successful UNLOAD statement preparation."""
        # Create Redshift config with UNLOAD enabled
        credentials = StandardCredentials(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
        )
        unload_config = UnloadConfig(
            s3_bucket="my-bucket",
            iam_role_arn="arn:aws:iam::123456789012:role/RedshiftUnloadRole",
            s3_prefix="exports",
            file_format="parquet",
            max_file_size_mb=256,
        )

        with patch(
            "shared.odbc.driver_utils.is_driver_available",
            return_value=True,
        ):
            redshift_config = RedshiftConnectionConfig(
                credentials=credentials,
                unload_config=unload_config,
            )

        source_connections_config = {"redshift": redshift_config}

        manager = _create_task_manager(source_connections_config)

        select_query = "SELECT * FROM users WHERE active = true"
        target_id = "workflows/1/tables/2/partition_0"

        unload_statement, data_source_type = manager._prepare_unload_statement(
            engine="redshift",
            select_query=select_query,
            target_id=target_id,
        )

        # Verify data source type
        assert data_source_type == "unload"

        # Verify UNLOAD statement components
        assert "UNLOAD" in unload_statement
        assert "SELECT * FROM users WHERE active = true" in unload_statement
        assert "s3://my-bucket/exports/workflows/1/tables/2/partition_0/" in unload_statement
        assert "IAM_ROLE 'arn:aws:iam::123456789012:role/RedshiftUnloadRole'" in unload_statement
        assert "FORMAT AS PARQUET" in unload_statement

    def test_prepare_unload_statement_rejects_non_redshift_engine(self):
        """Test that non-Redshift engines are rejected for S3_UNLOAD."""
        manager = _create_task_manager()

        with pytest.raises(ValueError) as exc_info:
            manager._prepare_unload_statement(
                engine="sqlserver",
                select_query="SELECT * FROM users",
                target_id="path/to/output",
            )

        assert "only supported for Redshift" in str(exc_info.value)

    def test_prepare_unload_statement_rejects_empty_target_id(self):
        """Test that empty target_id is rejected."""
        manager = _create_task_manager()

        with pytest.raises(ValueError) as exc_info:
            manager._prepare_unload_statement(
                engine="redshift",
                select_query="SELECT * FROM users",
                target_id="",
            )

        assert "requires 'target_id'" in str(exc_info.value)

    def test_prepare_unload_statement_rejects_none_target_id(self):
        """Test that None target_id is rejected."""
        manager = _create_task_manager()

        with pytest.raises(ValueError) as exc_info:
            manager._prepare_unload_statement(
                engine="redshift",
                select_query="SELECT * FROM users",
                target_id=None,
            )

        assert "requires 'target_id'" in str(exc_info.value)

    def test_prepare_unload_statement_rejects_missing_unload_config(self):
        """Test that missing UNLOAD configuration is rejected."""
        # Create Redshift config WITHOUT UNLOAD configuration
        credentials = StandardCredentials(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
        )

        with patch(
            "shared.odbc.driver_utils.is_driver_available",
            return_value=True,
        ):
            redshift_config = RedshiftConnectionConfig(
                credentials=credentials,
                unload_config=None,  # No UNLOAD config
            )

        source_connections_config = {"redshift": redshift_config}
        manager = _create_task_manager(source_connections_config)

        with pytest.raises(ValueError) as exc_info:
            manager._prepare_unload_statement(
                engine="redshift",
                select_query="SELECT * FROM users",
                target_id="path/to/output",
            )

        assert "UNLOAD configuration" in str(exc_info.value)


@pytest.mark.skipif(
    not TASK_MANAGER_AVAILABLE or not SQLSERVER_CONFIG_AVAILABLE or not REDSHIFT_CONFIG_AVAILABLE,
    reason="TaskManager or config classes not available (missing shared module)",
)
class TestTaskManagerPrepareUnloadStatementWithSQLServer:
    """Tests that require SQLServerConnectionConfig."""

    def test_prepare_unload_statement_rejects_non_redshift_connection_config(self):
        """Test that non-Redshift connection configs are rejected."""
        # Create SQL Server config for 'redshift' engine (misconfiguration)
        sqlserver_config = SQLServerConnectionConfig(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
        )

        source_connections_config = {"redshift": sqlserver_config}
        manager = _create_task_manager(source_connections_config)

        with pytest.raises(ValueError) as exc_info:
            manager._prepare_unload_statement(
                engine="redshift",
                select_query="SELECT * FROM users",
                target_id="path/to/output",
            )

        assert "Expected RedshiftConnectionConfig" in str(exc_info.value)


@pytest.mark.skipif(
    not TASK_MANAGER_AVAILABLE,
    reason="TaskManager not available (missing shared module)",
)
class TestTaskManagerResolveDataSourceType:
    """
    Test suite for TaskManager._resolve_data_source_type method.

    Validates the data source type resolution logic based on bulk flag,
    connection configuration, and statement analysis.
    """

    def test_resolve_returns_odbc_for_non_bulk_operation(self):
        """Test that ODBC is returned for non-bulk operations."""
        manager = _create_task_manager()

        result = manager._resolve_data_source_type(
            engine="sqlserver",
            is_bulk_operation=False,
        )

        assert result == "odbc"

    def test_resolve_returns_unload_for_redshift_unload_statement(self):
        """Test that UNLOAD type is returned for Redshift UNLOAD statements."""
        manager = _create_task_manager()

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

        result = manager._resolve_data_source_type(
            engine="redshift",
            is_bulk_operation=False,
            statement=unload_statement,
        )

        assert result == "unload"

    def test_resolve_returns_odbc_for_unload_on_non_redshift(self):
        """Test that ODBC is returned when UNLOAD is used with non-Redshift engine."""
        mock_logger = Mock()
        manager = _create_task_manager(mock_logger=mock_logger)

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

        result = manager._resolve_data_source_type(
            engine="sqlserver",
            is_bulk_operation=False,
            statement=unload_statement,
        )

        # Should fall back to ODBC and log a warning
        assert result == "odbc"
        mock_logger.warning.assert_called()
        warning_message = str(mock_logger.warning.call_args)
        assert "sqlserver" in warning_message

    def test_resolve_returns_odbc_for_select_statement(self):
        """Test that ODBC is returned for regular SELECT statements."""
        manager = _create_task_manager()

        result = manager._resolve_data_source_type(
            engine="redshift",
            is_bulk_operation=False,
            statement="SELECT * FROM users",
        )

        assert result == "odbc"

    def test_resolve_returns_odbc_for_engine_without_bulk_utility(self):
        """Test that ODBC is returned for engines without bulk utilities."""
        manager = _create_task_manager()

        result = manager._resolve_data_source_type(
            engine="postgresql",  # No bulk utility defined for PostgreSQL
            is_bulk_operation=True,
        )

        assert result == "odbc"

    def test_resolve_prioritizes_unload_over_bulk(self):
        """Test that UNLOAD statement detection takes priority over bulk flag."""
        manager = _create_task_manager()

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

        result = manager._resolve_data_source_type(
            engine="redshift",
            is_bulk_operation=True,  # Bulk flag is set but UNLOAD should take priority
            statement=unload_statement,
        )

        assert result == "unload"


@pytest.mark.skipif(
    not TASK_MANAGER_AVAILABLE or not SQLSERVER_CONFIG_AVAILABLE,
    reason="TaskManager or SQLServerConnectionConfig not available (missing shared module)",
)
class TestTaskManagerResolveDataSourceTypeWithSQLServer:
    """Tests that require SQLServerConnectionConfig for BCP testing."""

    def test_resolve_returns_bcp_for_sqlserver_bulk_operation(self):
        """Test that BCP is returned for SQL Server bulk operations with use_bcp=True."""
        sqlserver_config = SQLServerConnectionConfig(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
            use_bcp=True,
        )

        source_connections_config = {"sqlserver": sqlserver_config}
        manager = _create_task_manager(source_connections_config)

        result = manager._resolve_data_source_type(
            engine="sqlserver",
            is_bulk_operation=True,
        )

        assert result == "bcp"

    def test_resolve_returns_odbc_for_bulk_without_bcp_enabled(self):
        """Test that ODBC is returned for bulk operations when BCP is not enabled."""
        sqlserver_config = SQLServerConnectionConfig(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
            use_bcp=False,  # BCP disabled
        )

        source_connections_config = {"sqlserver": sqlserver_config}
        manager = _create_task_manager(source_connections_config)

        result = manager._resolve_data_source_type(
            engine="sqlserver",
            is_bulk_operation=True,
        )

        assert result == "odbc"


@pytest.mark.skipif(
    not TASK_MANAGER_AVAILABLE or not REDSHIFT_CONFIG_AVAILABLE,
    reason="TaskManager or Redshift config not available (missing shared module)",
)
class TestTaskManagerUnloadIntegration:
    """
    Integration tests for TaskManager UNLOAD functionality.

    Validates the complete flow of UNLOAD task processing.
    """

    def test_s3_unload_target_type_triggers_unload_preparation(self):
        """Test that s3:unload target type triggers UNLOAD statement preparation."""
        from data_exchange_agent.constants import task_keys

        # Create Redshift config with UNLOAD enabled
        credentials = StandardCredentials(
            database="testdb",
            username="testuser",
            host="localhost",
            password="testpass",
            auto_detect_driver=False,
        )
        unload_config = UnloadConfig(
            s3_bucket="my-bucket",
            iam_role_arn="arn:aws:iam::123:role/Role",
        )

        with patch(
            "shared.odbc.driver_utils.is_driver_available",
            return_value=True,
        ):
            redshift_config = RedshiftConnectionConfig(
                credentials=credentials,
                unload_config=unload_config,
            )

        source_connections_config = {"redshift": redshift_config}
        manager = _create_task_manager(source_connections_config)

        # Test task with s3:unload target type
        task = {
            task_keys.ID: "test-task-1",
            task_keys.ENGINE: "redshift",
            task_keys.STATEMENT: "SELECT * FROM users",
            task_keys.TARGET_TYPE: task_keys.TargetTypeValues.S3_UNLOAD,
            task_keys.TARGET_ID: "workflows/1/partition_0",
        }

        # Call _prepare_unload_statement through process_task logic
        unload_statement, data_source_type = manager._prepare_unload_statement(
            engine=task[task_keys.ENGINE],
            select_query=task[task_keys.STATEMENT],
            target_id=task[task_keys.TARGET_ID],
        )

        # Verify result
        assert data_source_type == "unload"
        assert "UNLOAD" in unload_statement
        assert "SELECT * FROM users" in unload_statement
