"""
Tests for the SnowflakeStoredProcedureTaskSourceAdapter class.

This module tests the SnowflakeStoredProcedureTaskSourceAdapter to ensure it properly
handles task retrieval, completion, and failure operations with Snowflake stored procedures.
"""

import json
import unittest

from unittest.mock import MagicMock, Mock

from data_exchange_agent import custom_exceptions
from data_exchange_agent.constants import config_keys
from data_exchange_agent.snowflake_metadata_locations import (
    qualified_data_migration_metadata_schema,
)
from data_exchange_agent.task_sources.snowflake_stored_procedure import (
    SnowflakeStoredProcedureTaskSourceAdapter,
)


class TestSnowflakeStoredProcedureTaskSourceAdapter(unittest.TestCase):
    """
    Comprehensive test suite for the SnowflakeStoredProcedureTaskSourceAdapter class.

    This test class validates the adapter's core functionality, including:
    - Initialization with proper configuration
    - Task retrieval from Snowflake stored procedures
    - Task completion status updates
    - Task failure status updates
    - Error handling for various failure scenarios
    - Proper interaction with SnowflakeDataSource

    Tests use extensive mocking to isolate the adapter from external
    dependencies like Snowflake connections, ensuring reliable and fast test execution.
    """

    def setUp(self):
        """
        Set up test fixtures before each test method.

        Creates a mock configuration and SnowflakeStoredProcedureTaskSourceAdapter instance
        to avoid external dependencies during testing.
        """
        # Create mock configuration
        self.mock_task_source = Mock()
        self.mock_task_source.connection_name = "test_snowflake_connection"

        self.mock_program_config = MagicMock()
        self.mock_program_config.__getitem__.side_effect = lambda key: {
            "task_source": self.mock_task_source,
            "application.agent_id": "test_agent_123",
        }[key]
        self.mock_program_config.get.side_effect = lambda key, default=None: default

        # Create mock snowflake datasource
        self.mock_sf_datasource = MagicMock()
        self.mock_sf_datasource.is_closed.return_value = False

    def _create_adapter(self, snowflake_datasource=None):
        """Create an adapter with mocked dependencies."""
        return SnowflakeStoredProcedureTaskSourceAdapter(
            program_config=self.mock_program_config,
            snowflake_datasource=snowflake_datasource or self.mock_sf_datasource,
        )

    def test_initialization_success(self):
        """Test successful initialization with valid configuration."""
        adapter = self._create_adapter()

        self.assertEqual(adapter.agent_id, "test_agent_123")
        self.assertEqual(adapter.snowflake_datasource, self.mock_sf_datasource)

    def test_initialization_missing_agent_id(self):
        """Test initialization fails when agent_id is missing."""
        mock_program_config = MagicMock()

        def mock_getitem(key):
            if key == "task_source":
                return self.mock_task_source
            raise KeyError(key)

        mock_program_config.__getitem__.side_effect = mock_getitem

        with self.assertRaises(custom_exceptions.ConfigurationError) as context:
            SnowflakeStoredProcedureTaskSourceAdapter(
                program_config=mock_program_config,
                snowflake_datasource=self.mock_sf_datasource,
            )

        self.assertIn("Agent ID is missing or incomplete", str(context.exception))

    def test_get_tasks_success(self):
        """Test successful task retrieval from Snowflake stored procedure."""
        # Mock raw task data returned from Snowflake
        raw_task_payload = {
            "source_type": "sqlserver",
            "database": "test_db",
            "schema": "dbo",
            "statement_location_id": "SELECT * FROM users",
            "target_type": "s3",
            "target_id": "s3://bucket/path/data",
        }

        raw_tasks = [
            {
                "ID": "task_001",
                "PAYLOAD": json.dumps(raw_task_payload),
            }
        ]

        self.mock_sf_datasource.execute_statement.return_value = raw_tasks

        adapter = self._create_adapter()
        result = adapter.get_tasks()

        # Verify the stored procedure call uses parameterized query
        call_args = self.mock_sf_datasource.execute_statement.call_args
        statement = call_args[0][0]
        params = call_args[0][1]
        qm = qualified_data_migration_metadata_schema(self.mock_program_config)
        self.assertIn(f"CALL {qm}.PULL_TASKS(%s, %s, %s, %s)", statement)
        self.assertEqual(params, ("test_agent_123", "data-exchange-agent", None, 1))

        # Verify the result
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["id"], "task_001")
        self.assertEqual(result[0]["engine"], "sqlserver")
        self.assertEqual(result[0]["database"], "test_db")
        self.assertEqual(result[0]["schema"], "dbo")
        self.assertEqual(result[0]["statement"], "SELECT * FROM users")
        self.assertEqual(result[0]["target_type"], "s3")
        self.assertEqual(result[0]["target_id"], "s3://bucket/path/data")

    def test_get_tasks_multiple_tasks(self):
        """Test retrieval of multiple tasks from Snowflake stored procedure."""
        # Mock multiple tasks
        raw_tasks = [
            {
                "ID": "task_001",
                "PAYLOAD": json.dumps(
                    {
                        "source_type": "sqlserver",
                        "database": "db1",
                        "schema": "dbo",
                        "statement_location_id": "SELECT 1",
                        "target_type": "s3",
                        "target_id": "s3://bucket1/path1",
                    }
                ),
            },
            {
                "ID": "task_002",
                "PAYLOAD": json.dumps(
                    {
                        "source_type": "sqlserver",
                        "database": "db2",
                        "schema": "dbo",
                        "statement_location_id": "SELECT 2",
                        "target_type": "azure",
                        "target_id": "azure://container/path2",
                    }
                ),
            },
        ]

        self.mock_sf_datasource.execute_statement.return_value = raw_tasks

        adapter = self._create_adapter()
        result = adapter.get_tasks()

        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]["id"], "task_001")
        self.assertEqual(result[0]["engine"], "sqlserver")
        self.assertEqual(result[1]["id"], "task_002")
        self.assertEqual(result[1]["engine"], "sqlserver")

    def test_get_tasks_empty_list(self):
        """Test get_tasks when no tasks are available."""
        self.mock_sf_datasource.execute_statement.return_value = []

        adapter = self._create_adapter()
        result = adapter.get_tasks()

        self.assertEqual(result, [])
        self.assertIsInstance(result, list)

    def test_get_tasks_execution_failure(self):
        """Test get_tasks when Snowflake execution fails."""
        self.mock_sf_datasource.execute_statement.side_effect = Exception(
            "Snowflake connection error"
        )

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.get_tasks()

        self.assertIn(
            "Failed to pull tasks executing statement", str(context.exception)
        )
        self.assertIn("PULL_TASKS", str(context.exception))
        # Verify connection was closed after error
        self.mock_sf_datasource.close_connection.assert_called_once()

    def test_get_tasks_json_parse_error(self):
        """Test get_tasks when JSON payload parsing fails."""
        # Return invalid JSON
        raw_tasks = [{"ID": "task_001", "PAYLOAD": "invalid json{"}]
        self.mock_sf_datasource.execute_statement.return_value = raw_tasks

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.get_tasks()

        self.assertIn(
            "Failed to pull tasks executing statement", str(context.exception)
        )

    def test_get_tasks_creates_connection_if_closed(self):
        """Test get_tasks creates connection if it's closed."""
        self.mock_sf_datasource.is_closed.return_value = True
        self.mock_sf_datasource.execute_statement.return_value = []

        adapter = self._create_adapter()
        adapter.get_tasks()

        self.mock_sf_datasource.create_connection.assert_called_once()

    def test_complete_task_success(self):
        """Test successful task completion."""
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"COMPLETE_TASK": True}]
        )

        adapter = self._create_adapter()
        adapter.complete_task("123")

        # Verify the stored procedure call uses parameterized query
        call_args = self.mock_sf_datasource.execute_statement.call_args
        statement = call_args[0][0]
        params = call_args[0][1]
        qm = qualified_data_migration_metadata_schema(self.mock_program_config)
        self.assertIn(f"CALL {qm}.COMPLETE_TASK(%s, %s)", statement)
        self.assertEqual(params, ("test_agent_123", 123))

    def test_complete_task_returns_false(self):
        """Test complete_task when stored procedure returns False."""
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"COMPLETE_TASK": False}]
        )

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.complete_task("123")

        self.assertIn("Failed to mark task '123' as completed", str(context.exception))

    def test_complete_task_returns_none(self):
        """Test complete_task when stored procedure returns None."""
        self.mock_sf_datasource.execute_statement.return_value = iter([])

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.complete_task("123")

        self.assertIn("Failed to mark task '123' as completed", str(context.exception))

    def test_complete_task_execution_failure(self):
        """Test complete_task when Snowflake execution fails."""
        self.mock_sf_datasource.execute_statement.side_effect = Exception(
            "Database error"
        )

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.complete_task("123")

        self.assertIn("Failed to mark task '123' as completed", str(context.exception))
        self.assertIn("COMPLETE_TASK", str(context.exception))
        # Verify connection was closed after error
        self.mock_sf_datasource.close_connection.assert_called_once()

    def test_fail_task_success_without_error_message(self):
        """Test successful task failure marking without error message."""
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"FAIL_TASK": True}]
        )

        adapter = self._create_adapter()
        adapter.fail_task("456")

        # Verify the stored procedure call uses parameterized query
        call_args = self.mock_sf_datasource.execute_statement.call_args
        statement = call_args[0][0]
        params = call_args[0][1]
        qm = qualified_data_migration_metadata_schema(self.mock_program_config)
        self.assertIn(f"CALL {qm}.FAIL_TASK(%s, %s, %s)", statement)
        self.assertEqual(params, (456, "test_agent_123", "No error message provided."))

    def test_fail_task_success_with_error_message(self):
        """Test successful task failure marking with error message."""
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"FAIL_TASK": True}]
        )

        adapter = self._create_adapter()
        adapter.fail_task("456", "Connection timeout error")

        # Verify the stored procedure call uses parameterized query
        call_args = self.mock_sf_datasource.execute_statement.call_args
        statement = call_args[0][0]
        params = call_args[0][1]
        qm = qualified_data_migration_metadata_schema(self.mock_program_config)
        self.assertIn(f"CALL {qm}.FAIL_TASK(%s, %s, %s)", statement)
        self.assertEqual(params, (456, "test_agent_123", "Connection timeout error"))

    def test_fail_task_with_error_message_containing_quotes(self):
        """
        Test fail_task with error message containing single quotes.

        With parameterized queries, the driver handles escaping so the raw
        message is passed as-is in the params tuple.
        """
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"FAIL_TASK": True}]
        )

        adapter = self._create_adapter()
        adapter.fail_task("456", "Error: Can't connect to 'database'")

        # Verify the raw (unescaped) message is passed via bind parameters
        call_args = self.mock_sf_datasource.execute_statement.call_args
        params = call_args[0][1]
        self.assertEqual(params[2], "Error: Can't connect to 'database'")

    def test_fail_task_returns_false(self):
        """Test fail_task when stored procedure returns False."""
        self.mock_sf_datasource.execute_statement.return_value = iter(
            [{"FAIL_TASK": False}]
        )

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.fail_task("456", "Test error")

        self.assertIn("Failed to mark task '456' as failed", str(context.exception))

    def test_fail_task_returns_none(self):
        """Test fail_task when stored procedure returns None."""
        self.mock_sf_datasource.execute_statement.return_value = iter([])

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.fail_task("456")

        self.assertIn("Failed to mark task '456' as failed", str(context.exception))

    def test_fail_task_execution_failure(self):
        """Test fail_task when Snowflake execution fails."""
        self.mock_sf_datasource.execute_statement.side_effect = Exception(
            "Network error"
        )

        adapter = self._create_adapter()

        with self.assertRaises(Exception) as context:
            adapter.fail_task("456", "Original error")

        self.assertIn("Failed to mark task '456' as failed", str(context.exception))
        self.assertIn("FAIL_TASK", str(context.exception))
        # Verify connection was closed after error
        self.mock_sf_datasource.close_connection.assert_called_once()

    def test_adapter_constants(self):
        """Test that adapter class constants are properly defined."""
        self.assertEqual(
            SnowflakeStoredProcedureTaskSourceAdapter.AGENT_TYPE, "data-exchange-agent"
        )
        self.assertEqual(
            SnowflakeStoredProcedureTaskSourceAdapter.MAX_TASKS_PER_FETCH, 1
        )
        self.assertEqual(
            SnowflakeStoredProcedureTaskSourceAdapter.PULL_TASKS_SP_NAME, "PULL_TASKS"
        )
        self.assertEqual(
            SnowflakeStoredProcedureTaskSourceAdapter.COMPLETE_TASK_SP_NAME,
            "COMPLETE_TASK",
        )
        self.assertEqual(
            SnowflakeStoredProcedureTaskSourceAdapter.FAIL_TASK_SP_NAME, "FAIL_TASK"
        )

    def test_custom_snowflake_metadata_from_config(self):
        """Task-queue CALL prefix uses application.snowflake_*_metadata when set."""
        custom_config = MagicMock()
        custom_config.__getitem__.side_effect = lambda key: {
            config_keys.APPLICATION__AGENT_ID: "test_agent_123",
        }[key]

        def custom_get(key, default=None):
            if key == config_keys.APPLICATION__SNOWFLAKE_DATABASE_FOR_METADATA:
                return "MYDB"
            if (
                key
                == config_keys.APPLICATION__SNOWFLAKE_SCHEMA_FOR_DATA_MIGRATION_METADATA
            ):
                return "MYSCHEMA"
            return default

        custom_config.get.side_effect = custom_get

        adapter = SnowflakeStoredProcedureTaskSourceAdapter(
            program_config=custom_config,
            snowflake_datasource=self.mock_sf_datasource,
        )
        self.assertEqual(
            adapter._qualified_data_migration_metadata_schema, "MYDB.MYSCHEMA"
        )

    def test_adapter_inherits_from_task_source_adapter(self):
        """Test that SnowflakeStoredProcedureTaskSourceAdapter inherits from TaskSourceAdapter."""
        from data_exchange_agent.interfaces.task_source_adapter import (
            TaskSourceAdapter,
        )

        adapter = self._create_adapter()
        self.assertIsInstance(adapter, TaskSourceAdapter)

    def test_connection_reused_across_calls(self):
        """Test that the Snowflake connection is reused across multiple calls."""
        self.mock_sf_datasource.execute_statement.return_value = []

        adapter = self._create_adapter()

        # Make multiple calls
        adapter.get_tasks()
        adapter.get_tasks()
        adapter.get_tasks()

        # Connection should not be recreated since it's not closed
        self.mock_sf_datasource.create_connection.assert_not_called()

    def test_get_tasks_with_complex_statement(self):
        """Test get_tasks with complex SQL statement in payload."""
        complex_statement = """
            WITH cte AS (
                SELECT * FROM table1
                WHERE condition = 'test'
            )
            SELECT a.*, b.* FROM cte a
            JOIN table2 b ON a.id = b.id
        """

        raw_task_payload = {
            "source_type": "sqlserver",
            "database": "test_db",
            "schema": "dbo",
            "statement_location_id": complex_statement,
            "target_type": "s3",
            "target_id": "s3://bucket/path",
        }

        raw_tasks = [{"ID": "task_complex", "PAYLOAD": json.dumps(raw_task_payload)}]
        self.mock_sf_datasource.execute_statement.return_value = raw_tasks

        adapter = self._create_adapter()
        result = adapter.get_tasks()

        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["statement"], complex_statement)


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