import unittest

from unittest.mock import MagicMock, patch, call

from data_exchange_agent.data_sources.sf_connection import SnowflakeDataSource
from data_exchange_agent.uploaders.sf_stage_uploader import SFStageUploader


class TestSFStageUploader(unittest.TestCase):
    """
    Comprehensive test suite for the SFStageUploader class.

    This test class validates the SFStageUploader's functionality, including:
    - Connection establishment to Snowflake data sources
    - File upload operations to Snowflake stages
    - Stage refresh operations after uploads
    - Error handling for upload failures and connection issues
    - Proper cleanup and disconnection procedures
    - Integration with dependency injection

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

    def setUp(self):
        """Set up test fixtures before each test method."""
        self.uploader = SFStageUploader(cloud_storage_toml={})
        self.mock_snowflake_datasource = MagicMock(spec=SnowflakeDataSource)
        self.mock_snowflake_datasource.__enter__.return_value = self.mock_snowflake_datasource
        self.mock_snowflake_datasource.__exit__.return_value = False
        self.uploader.snowflake_datasource = self.mock_snowflake_datasource

    def test_init(self):
        """Test SFStageUploader initialization."""
        self.assertIsInstance(self.uploader, SFStageUploader)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_connect_with_existing_open_connection(self, mock_provide):
        """Test connect when Snowflake connection already exists and is open."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        self.uploader.connect()

        # Should not call create_connection since connection is already open
        self.mock_snowflake_datasource.create_connection.assert_not_called()
        self.mock_snowflake_datasource.is_closed.assert_called_once()

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_connect_with_closed_connection(self, mock_provide):
        """Test connect when Snowflake connection is closed."""
        self.mock_snowflake_datasource.is_closed.return_value = True

        self.uploader.connect()

        # Should call create_connection since connection is closed
        self.mock_snowflake_datasource.is_closed.assert_called_once()
        self.mock_snowflake_datasource.create_connection.assert_called_once()

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_connect_with_no_existing_connection(self, mock_provide):
        """Test connect when no Snowflake connection exists."""
        # Simulate no existing connection by making the datasource None initially
        self.mock_snowflake_datasource.is_closed.return_value = True

        self.uploader.connect()

        # Should call create_connection
        self.mock_snowflake_datasource.create_connection.assert_called_once()

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_disconnect_with_open_connection(self, mock_provide):
        """Test disconnect when Snowflake connection is open."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        self.uploader.disconnect()

        # Should call close_connection since connection is open
        self.mock_snowflake_datasource.is_closed.assert_called_once()
        self.mock_snowflake_datasource.close_connection.assert_called_once()

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_disconnect_with_closed_connection(self, mock_provide):
        """Test disconnect when Snowflake connection is already closed."""
        self.mock_snowflake_datasource.is_closed.return_value = True

        self.uploader.disconnect()

        # Should not call close_connection since connection is already closed
        self.mock_snowflake_datasource.is_closed.assert_called_once()
        self.mock_snowflake_datasource.close_connection.assert_not_called()

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_success(self, mock_provide):
        """Test successful file upload to Snowflake stage with stage refresh."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        # Mock successful upload response for PUT command, then REFRESH
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []  # ALTER STAGE REFRESH returns empty result
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file(
            "/path/to/test.parquet",
            "@my_stage/destination/",
        )

        # Verify both PUT and ALTER STAGE REFRESH commands were executed
        expected_put_command = "PUT file:///path/to/test.parquet @my_stage/destination/ OVERWRITE = TRUE"
        expected_refresh_command = "ALTER STAGE my_stage REFRESH SUBPATH = 'destination/'"
        self.assertEqual(self.mock_snowflake_datasource.execute_statement.call_count, 2)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put_command)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_not_found(self, mock_provide):
        """Test upload with non-existent file - Snowflake handles file validation."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        # Mock failed upload response (file not found)
        upload_result = [{"status": "ERROR"}]
        self.mock_snowflake_datasource.execute_statement.return_value = iter(upload_result)

        with self.assertRaises(Exception) as context:
            self.uploader.upload_file(
                "/path/to/nonexistent.parquet",
                "@my_stage/destination/",
            )

        self.assertIn(
            "Failed to upload file /path/to/nonexistent.parquet to @my_stage/destination/",
            str(context.exception),
        )

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_upload_failed(self, mock_provide):
        """Test upload failure handling."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        # Mock failed upload response (status is not "UPLOADED")
        upload_result = [{"status": "FAILED"}]
        self.mock_snowflake_datasource.execute_statement.return_value = iter(upload_result)

        with self.assertRaises(Exception) as context:
            self.uploader.upload_file(
                "/path/to/test.parquet",
                "@my_stage/destination/",
            )

        self.assertIn(
            "Failed to upload file /path/to/test.parquet to @my_stage/destination/",
            str(context.exception),
        )

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_no_response(self, mock_provide):
        """Test upload with no response from Snowflake."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        # Mock empty response
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        with self.assertRaises(Exception) as context:
            self.uploader.upload_file(
                "/path/to/test.parquet",
                "@my_stage/destination/",
            )

        self.assertIn(
            "Failed to upload file /path/to/test.parquet to @my_stage/destination/",
            str(context.exception),
        )

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_snowflake_exception(self, mock_provide):
        """Test upload_file handles Snowflake exceptions."""
        self.mock_snowflake_datasource.is_closed.return_value = False

        # Mock Snowflake to raise exception
        self.mock_snowflake_datasource.execute_statement.side_effect = Exception("Snowflake error")

        with self.assertRaises(Exception) as context:
            self.uploader.upload_file(
                "/path/to/test.parquet",
                "@my_stage/destination/",
            )

        self.assertIn("Snowflake error", str(context.exception))


class TestSFStageUploaderRefreshStage(unittest.TestCase):
    """Test suite for the _refresh_stage method."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        self.uploader = SFStageUploader(cloud_storage_toml={})
        self.mock_snowflake_datasource = MagicMock(spec=SnowflakeDataSource)
        self.mock_snowflake_datasource.__enter__.return_value = self.mock_snowflake_datasource
        self.mock_snowflake_datasource.__exit__.return_value = False
        self.uploader.snowflake_datasource = self.mock_snowflake_datasource

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_with_subpath(self, mock_provide):
        """Test stage refresh with a subpath."""
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        self.uploader._refresh_stage(self.mock_snowflake_datasource, "@my_stage/folder/subfolder/")

        expected_command = "ALTER STAGE my_stage REFRESH SUBPATH = 'folder/subfolder/'"
        self.mock_snowflake_datasource.execute_statement.assert_called_once_with(expected_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_without_subpath(self, mock_provide):
        """Test stage refresh without a subpath (stage root only)."""
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        self.uploader._refresh_stage(self.mock_snowflake_datasource, "@my_stage")

        expected_command = "ALTER STAGE my_stage REFRESH"
        self.mock_snowflake_datasource.execute_statement.assert_called_once_with(expected_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_with_trailing_slash_only(self, mock_provide):
        """Test stage refresh with only trailing slash (no subpath)."""
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        self.uploader._refresh_stage(self.mock_snowflake_datasource, "@my_stage/")

        expected_command = "ALTER STAGE my_stage REFRESH"
        self.mock_snowflake_datasource.execute_statement.assert_called_once_with(expected_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_with_leading_slash(self, mock_provide):
        """Test stage refresh with leading slash in path."""
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        self.uploader._refresh_stage(self.mock_snowflake_datasource, "/@my_stage/folder/")

        expected_command = "ALTER STAGE my_stage REFRESH SUBPATH = 'folder/'"
        self.mock_snowflake_datasource.execute_statement.assert_called_once_with(expected_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_with_qualified_stage_name(self, mock_provide):
        """Test stage refresh with fully qualified stage name (database.schema.stage)."""
        self.mock_snowflake_datasource.execute_statement.return_value = iter([])

        self.uploader._refresh_stage(self.mock_snowflake_datasource, "@database.schema.stage/path/")

        expected_command = "ALTER STAGE database.schema.stage REFRESH SUBPATH = 'path/'"
        self.mock_snowflake_datasource.execute_statement.assert_called_once_with(expected_command)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_error_raises_exception(self, mock_provide):
        """Test that refresh stage error raises exception with proper message."""
        self.mock_snowflake_datasource.execute_statement.side_effect = Exception("Permission denied")

        with self.assertRaises(Exception) as context:
            self.uploader._refresh_stage(self.mock_snowflake_datasource, "@my_stage/folder/")

        self.assertIn("Failed to refresh stage my_stage with subpath 'folder/'", str(context.exception))
        self.assertIn("Permission denied", str(context.exception))

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_refresh_stage_error_without_subpath_raises_exception(self, mock_provide):
        """Test that refresh stage error raises exception without subpath in message."""
        self.mock_snowflake_datasource.execute_statement.side_effect = Exception("Stage not found")

        with self.assertRaises(Exception) as context:
            self.uploader._refresh_stage(self.mock_snowflake_datasource, "@my_stage")

        self.assertIn("Failed to refresh stage my_stage.", str(context.exception))
        self.assertIn("Stage not found", str(context.exception))


class TestSFStageUploaderUploadFiles(unittest.TestCase):
    """Test suite for the upload_files method."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        self.uploader = SFStageUploader(cloud_storage_toml={})
        self.mock_snowflake_datasource = MagicMock(spec=SnowflakeDataSource)
        self.mock_snowflake_datasource.__enter__.return_value = self.mock_snowflake_datasource
        self.mock_snowflake_datasource.__exit__.return_value = False
        self.uploader.snowflake_datasource = self.mock_snowflake_datasource

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_files_success(self, mock_provide):
        """Test successful upload of multiple files with single stage refresh."""
        # Mock successful upload responses for each file, then one refresh
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),  # First file upload
            iter(upload_result),  # Second file upload
            iter(upload_result),  # Third file upload
            iter(refresh_result),  # Single refresh at the end
        ]

        self.uploader.upload_files(
            "/path/to/file1.parquet",
            "/path/to/file2.parquet",
            "/path/to/file3.parquet",
            destination_path="@my_stage/folder/",
        )

        # Should execute 3 PUT commands + 1 REFRESH command = 4 total
        self.assertEqual(self.mock_snowflake_datasource.execute_statement.call_count, 4)

        # Verify PUT commands
        expected_calls = [
            call("PUT file:///path/to/file1.parquet @my_stage/folder/ OVERWRITE = TRUE"),
            call("PUT file:///path/to/file2.parquet @my_stage/folder/ OVERWRITE = TRUE"),
            call("PUT file:///path/to/file3.parquet @my_stage/folder/ OVERWRITE = TRUE"),
            call("ALTER STAGE my_stage REFRESH SUBPATH = 'folder/'"),
        ]
        self.mock_snowflake_datasource.execute_statement.assert_has_calls(expected_calls, any_order=False)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_files_single_file(self, mock_provide):
        """Test upload_files with a single file."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_files(
            "/path/to/single.parquet",
            destination_path="@my_stage/",
        )

        self.assertEqual(self.mock_snowflake_datasource.execute_statement.call_count, 2)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_files_fails_on_first_file(self, mock_provide):
        """Test that upload_files fails when first file upload fails."""
        upload_result = [{"status": "FAILED"}]
        self.mock_snowflake_datasource.execute_statement.return_value = iter(upload_result)

        with self.assertRaises(Exception) as context:
            self.uploader.upload_files(
                "/path/to/file1.parquet",
                "/path/to/file2.parquet",
                destination_path="@my_stage/folder/",
            )

        self.assertIn("Failed to upload file /path/to/file1.parquet", str(context.exception))
        # Only the first PUT should have been attempted
        self.assertEqual(self.mock_snowflake_datasource.execute_statement.call_count, 1)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_files_refresh_fails(self, mock_provide):
        """Test that upload_files fails when stage refresh fails after successful uploads."""
        upload_result = [{"status": "UPLOADED"}]
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),  # First file upload succeeds
            iter(upload_result),  # Second file upload succeeds
            Exception("Refresh failed"),  # Refresh fails
        ]

        with self.assertRaises(Exception) as context:
            self.uploader.upload_files(
                "/path/to/file1.parquet",
                "/path/to/file2.parquet",
                destination_path="@my_stage/folder/",
            )

        self.assertIn("Failed to refresh stage my_stage", str(context.exception))


class TestSFStageUploaderDifferentStageFormats(unittest.TestCase):
    """Test suite for different stage path formats."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        self.uploader = SFStageUploader(cloud_storage_toml={})
        self.mock_snowflake_datasource = MagicMock(spec=SnowflakeDataSource)
        self.mock_snowflake_datasource.__enter__.return_value = self.mock_snowflake_datasource
        self.mock_snowflake_datasource.__exit__.return_value = False
        self.uploader.snowflake_datasource = self.mock_snowflake_datasource

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_stage_only(self, mock_provide):
        """Test upload with stage name only (no subpath)."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file("/path/to/test.parquet", "@my_stage")

        expected_put = "PUT file:///path/to/test.parquet @my_stage OVERWRITE = TRUE"
        expected_refresh = "ALTER STAGE my_stage REFRESH"
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_stage_with_trailing_slash(self, mock_provide):
        """Test upload with stage name and trailing slash."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file("/path/to/test.parquet", "@my_stage/")

        expected_put = "PUT file:///path/to/test.parquet @my_stage/ OVERWRITE = TRUE"
        expected_refresh = "ALTER STAGE my_stage REFRESH"
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_stage_with_single_folder(self, mock_provide):
        """Test upload with single folder subpath."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file("/path/to/test.parquet", "@my_stage/folder/")

        expected_put = "PUT file:///path/to/test.parquet @my_stage/folder/ OVERWRITE = TRUE"
        expected_refresh = "ALTER STAGE my_stage REFRESH SUBPATH = 'folder/'"
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_stage_with_nested_folders(self, mock_provide):
        """Test upload with nested folder subpath."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file("/path/to/test.parquet", "@my_stage/folder/subfolder/deep/")

        expected_put = "PUT file:///path/to/test.parquet @my_stage/folder/subfolder/deep/ OVERWRITE = TRUE"
        expected_refresh = "ALTER STAGE my_stage REFRESH SUBPATH = 'folder/subfolder/deep/'"
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh)

    @patch("data_exchange_agent.uploaders.sf_stage_uploader.Provide")
    def test_upload_file_fully_qualified_stage(self, mock_provide):
        """Test upload with fully qualified stage name (database.schema.stage)."""
        upload_result = [{"status": "UPLOADED"}]
        refresh_result = []
        self.mock_snowflake_datasource.execute_statement.side_effect = [
            iter(upload_result),
            iter(refresh_result),
        ]

        self.uploader.upload_file("/path/to/test.parquet", "@database.schema.my_stage/folder/")

        expected_put = "PUT file:///path/to/test.parquet @database.schema.my_stage/folder/ OVERWRITE = TRUE"
        expected_refresh = "ALTER STAGE database.schema.my_stage REFRESH SUBPATH = 'folder/'"
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_put)
        self.mock_snowflake_datasource.execute_statement.assert_any_call(expected_refresh)


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