"""
Comprehensive test suite for the BCPDataSource class.

This test class validates the core functionality of the BCPDataSource class,
including initialization, data export, BCP command building, and error handling.
"""

import os
import tempfile
import unittest

from unittest.mock import Mock, patch

from data_exchange_agent.config.sections.connections.odbc.sqlserver import (
    SQLServerConnectionConfig,
)
from data_exchange_agent.data_sources.base import BaseDataSource
from data_exchange_agent.data_sources.bcp_data_source import BCPDataSource


class TestBCPDataSource(unittest.TestCase):
    """
    Comprehensive test suite for the BCPDataSource class.

    This test class validates the core functionality and ensures proper
    behavior under various conditions including normal operation,
    error scenarios, and edge cases.

    Tests use mocking where appropriate to isolate the component
    under test and ensure reliable, fast test execution.
    """

    def setUp(self):
        """Set up test fixtures before each test method."""
        self.engine = "sqlserver_test"
        self.statement = "SELECT * FROM users"
        self.temp_dir = tempfile.mkdtemp()
        self.mock_logger = Mock()

        # Mock connection config
        self.mock_connection_config = SQLServerConnectionConfig(
            username="testuser",
            password="testpass",
            database="testdb",
            host="localhost",
            port=1433,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        # Mock program config
        self.mock_program_config = Mock()
        self.mock_program_config.__getitem__ = Mock(side_effect=self._mock_config_getitem)

    def _mock_config_getitem(self, key):
        """Mock the config manager's __getitem__ method."""
        if key == "connections.source":
            return {self.engine: self.mock_connection_config}
        return None

    def tearDown(self):
        """Clean up after each test method."""
        # Clean up temp directory
        import shutil

        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)

    def test_bcp_data_source_is_base_data_source(self):
        """Test that BCPDataSource is a subclass of BaseDataSource."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        self.assertIsInstance(data_source, BaseDataSource)

    def test_initialization_with_defaults(self):
        """Test BCPDataSource initialization with default parameters."""
        with patch("data_exchange_agent.data_sources.bcp_data_source.build_actual_results_folder_path") as mock_build_path:
            mock_build_path.return_value = self.temp_dir

            data_source = BCPDataSource(
                engine=self.engine,
                statement=self.statement,
                logger=self.mock_logger,
                program_config=self.mock_program_config,
            )

            self.assertEqual(data_source.statement, self.statement)
            self.assertEqual(data_source.results_folder_path, self.temp_dir)
            self.assertEqual(data_source.base_file_name, "result")

    def test_initialization_with_custom_parameters(self):
        """Test BCPDataSource initialization with custom parameters."""
        custom_path = "/custom/results/path"
        custom_base_name = "custom_result"

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=custom_path,
            base_file_name=custom_base_name,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        self.assertEqual(data_source.statement, self.statement)
        self.assertEqual(data_source.results_folder_path, custom_path)
        self.assertEqual(data_source.base_file_name, custom_base_name)

    def test_build_bcp_command_with_query(self):
        """Test building BCP command with a SELECT query."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement="SELECT * FROM users WHERE id > 10",
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        self.assertIn("queryout", command)
        self.assertIn("SELECT * FROM users WHERE id > 10", command)
        self.assertIn("-U", command)
        self.assertIn("testuser", command)
        self.assertIn("-P", command)
        self.assertIn("testpass", command)
        self.assertIn("-d", command)
        self.assertIn("testdb", command)
        # Default delimiter (unit separator \x1f) and encoding
        self.assertIn("-t", command)
        self.assertIn("\x1f", command)
        self.assertIn("-r", command)
        self.assertIn("\\n", command)
        self.assertIn("-C", command)
        self.assertIn("UTF8", command)

    def test_build_bcp_command_with_table_name(self):
        """Test building BCP command with a table name."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement="dbo.users",
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        self.assertIn("out", command)
        self.assertIn("dbo.users", command)
        self.assertNotIn("queryout", command)

    def test_build_bcp_command_with_windows_auth_from_connection(self):
        """Test building BCP command with Windows Auth from connection config."""
        # Create connection config with Windows Auth
        windows_auth_config = SQLServerConnectionConfig(
            database="testdb",
            host="localhost",
            port=1433,
            use_windows_auth=True,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        # Create a new mock config that returns the Windows Auth config
        def mock_config_with_windows_auth(key):
            if key == "connections.source":
                return {self.engine: windows_auth_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_with_windows_auth)

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        self.assertIn("-T", command)
        self.assertNotIn("-U", command)
        self.assertNotIn("-P", command)

    def test_build_server_string_with_encryption(self):
        """Test building server string with encryption properties from connection config."""
        # Create connection config with encrypt=True
        encrypted_connection_config = SQLServerConnectionConfig(
            username="testuser",
            password="testpass",
            database="testdb",
            host="localhost",
            port=1433,
            encrypt=True,
            auto_detect_driver=False,
        )

        def mock_config_with_encryption(key):
            if key == "connections.source":
                return {self.engine: encrypted_connection_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_with_encryption)

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        server_string = data_source._build_server_string()

        self.assertIn("localhost", server_string)
        self.assertIn("Encrypt=yes", server_string)

    def test_build_server_string_with_trust_server_certificate(self):
        """Test building server string with trust server certificate."""
        self.mock_connection_config.extra_options = {"trustServerCertificate": True}

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        server_string = data_source._build_server_string()

        self.assertIn("TrustServerCertificate=yes", server_string)

    def test_build_server_string_without_properties(self):
        """Test building server string without encryption properties."""
        self.mock_connection_config.extra_options = {}

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        server_string = data_source._build_server_string()

        self.assertEqual(server_string, "localhost")

    def test_mask_password_in_command(self):
        """Test masking password in command string."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        command = [
            "bcp",
            "queryout",
            "-U",
            "testuser",
            "-P",
            "testpass",
            "-d",
            "testdb",
        ]
        masked_command = data_source._mask_password_in_command(command)

        self.assertIn("***", masked_command)
        self.assertNotIn("testpass", masked_command)

    def test_mask_password_with_windows_auth(self):
        """Test that password is not masked when using Windows Auth."""
        # Create connection config with Windows Auth
        windows_auth_config = SQLServerConnectionConfig(
            database="testdb",
            host="localhost",
            port=1433,
            use_windows_auth=True,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        # Create a new mock config that returns the Windows Auth config
        def mock_config_with_windows_auth(key):
            if key == "connections.source":
                return {self.engine: windows_auth_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_with_windows_auth)

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        command = ["bcp", "queryout", "-T", "-d", "testdb"]
        masked_command = data_source._mask_password_in_command(command)

        self.assertEqual(" ".join(command), masked_command)

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.Popen")
    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_success(self, mock_subprocess_run, mock_popen):
        """Test successful data export."""
        # Mock BCP version check
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 0
        mock_subprocess_run.return_value = mock_bcp_check

        # Mock BCP execution with Popen
        mock_process = Mock()
        mock_process.stdout = iter(["1000 rows copied\n"])
        mock_process.stderr = iter([])
        mock_process.wait.return_value = 0
        mock_popen.return_value = mock_process

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        result = data_source.export_data()

        self.assertTrue(result)
        self.assertEqual(mock_subprocess_run.call_count, 1)
        self.mock_logger.info.assert_called()

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_bcp_not_installed(self, mock_subprocess_run):
        """Test export_data when BCP is not installed."""
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 1
        mock_subprocess_run.return_value = mock_bcp_check

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

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

        self.assertIn("BCP utility is not installed", str(context.exception))

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.Popen")
    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_bcp_command_fails(self, mock_subprocess_run, mock_popen):
        """Test export_data when BCP command fails."""
        # Mock BCP version check
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 0
        mock_subprocess_run.return_value = mock_bcp_check

        # Mock BCP execution failure with Popen
        mock_process = Mock()
        mock_process.stdout = iter([])
        mock_process.stderr = iter(["Error connecting to database\n"])
        mock_process.wait.return_value = 1
        mock_popen.return_value = mock_process

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

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

        self.assertIn("BCP command failed", str(context.exception))
        self.assertIn("Error connecting to database", str(context.exception))

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.Popen")
    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_creates_output_directory(self, mock_subprocess_run, mock_popen):
        """Test that export_data creates the output directory if it doesn't exist."""
        # Mock BCP version check
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 0
        mock_subprocess_run.return_value = mock_bcp_check

        # Mock BCP execution with Popen
        mock_process = Mock()
        mock_process.stdout = iter(["Success\n"])
        mock_process.stderr = iter([])
        mock_process.wait.return_value = 0
        mock_popen.return_value = mock_process

        # Use a path that doesn't exist
        non_existent_path = os.path.join(self.temp_dir, "new_directory")

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=non_existent_path,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        data_source.export_data()

        self.assertTrue(os.path.exists(non_existent_path))

    def test_statement_property(self):
        """Test statement property returns the correct value."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        self.assertEqual(data_source.statement, self.statement)

    def test_results_folder_path_property(self):
        """Test results_folder_path property returns the correct value."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        self.assertEqual(data_source.results_folder_path, self.temp_dir)

    def test_base_file_name_property(self):
        """Test base_file_name property returns the correct value."""
        custom_base_name = "my_result"

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            base_file_name=custom_base_name,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        self.assertEqual(data_source.base_file_name, custom_base_name)

    def test_build_bcp_command_with_exec_statement(self):
        """Test building BCP command with EXEC statement."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement="EXEC sp_get_users",
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        self.assertIn("queryout", command)
        self.assertIn("EXEC sp_get_users", command)

    def test_build_bcp_command_with_cte(self):
        """Test building BCP command with CTE (WITH clause)."""
        statement = "WITH cte AS (SELECT * FROM users) SELECT * FROM cte"
        data_source = BCPDataSource(
            engine=self.engine,
            statement=statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        self.assertIn("queryout", command)
        self.assertIn(statement, command)

    def test_bcp_uses_default_delimiter(self):
        """Test that BCP always uses the unit separator delimiter."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Find the delimiter in command (follows -t flag)
        t_index = command.index("-t")
        delimiter = command[t_index + 1]
        self.assertEqual(delimiter, "\x1f")

    def test_bcp_uses_default_row_terminator(self):
        """Test that BCP always uses the default newline row terminator."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Find the row terminator in command (follows -r flag)
        r_index = command.index("-r")
        row_terminator = command[r_index + 1]
        self.assertEqual(row_terminator, "\\n")

    def test_bcp_uses_utf8_encoding(self):
        """Test that BCP always uses UTF8 encoding."""
        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Find the encoding in command (follows -C flag)
        c_index = command.index("-C")
        encoding = command[c_index + 1]
        self.assertEqual(encoding, "UTF8")

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.Popen")
    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_logs_masked_password(self, mock_subprocess_run, mock_popen):
        """Test that export_data logs command with masked password."""
        # Mock BCP version check
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 0
        mock_subprocess_run.return_value = mock_bcp_check

        # Mock BCP execution with Popen
        mock_process = Mock()
        mock_process.stdout = iter(["Success\n"])
        mock_process.stderr = iter([])
        mock_process.wait.return_value = 0
        mock_popen.return_value = mock_process

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        data_source.export_data()

        # Check that password was masked in log
        log_calls = [str(call) for call in self.mock_logger.info.call_args_list]
        password_in_log = any("testpass" in str(call) for call in log_calls)
        self.assertFalse(password_in_log)

    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.Popen")
    @patch("data_exchange_agent.data_sources.bcp_data_source.subprocess.run")
    def test_export_data_exception_handling(self, mock_subprocess_run, mock_popen):
        """Test that export_data handles exceptions properly."""
        # Mock BCP version check to succeed
        mock_bcp_check = Mock()
        mock_bcp_check.returncode = 0
        mock_subprocess_run.return_value = mock_bcp_check

        # Then raise an exception during actual BCP execution
        mock_popen.side_effect = Exception("Unexpected error")

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

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

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

    def test_build_server_string_with_encrypt_false(self):
        """Test building server string with encrypt=False."""
        # Create connection config with encrypt=False
        no_encrypt_config = SQLServerConnectionConfig(
            username="testuser",
            password="testpass",
            database="testdb",
            host="localhost",
            port=1433,
            encrypt=False,
            auto_detect_driver=False,
        )

        def mock_config_no_encrypt(key):
            if key == "connections.source":
                return {self.engine: no_encrypt_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_no_encrypt)

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        server_string = data_source._build_server_string()

        self.assertIn("Encrypt=no", server_string)

    def test_sql_auth_from_connection_config(self):
        """Test BCP uses SQL Auth when connection has use_windows_auth=False."""
        # Create connection config with SQL Server Auth
        self.mock_connection_config = SQLServerConnectionConfig(
            username="testuser",
            password="testpass",
            database="testdb",
            host="localhost",
            port=1433,
            use_windows_auth=False,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=self.mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Should use -U and -P flags, not -T
        self.assertIn("-U", command)
        self.assertIn("testuser", command)
        self.assertIn("-P", command)
        self.assertIn("testpass", command)
        self.assertNotIn("-T", command)

    def test_windows_auth_from_connection(self):
        """Test BCP uses Windows Auth from connection config."""
        # Create connection config with Windows Auth
        windows_auth_config = SQLServerConnectionConfig(
            database="testdb",
            host="localhost",
            port=1433,
            use_windows_auth=True,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        # Mock program config
        def mock_config_getitem(key):
            if key == "connections.source":
                return {self.engine: windows_auth_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_getitem)

        data_source = BCPDataSource(
            engine=self.engine,
            statement=self.statement,
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Should use -T flag
        self.assertIn("-T", command)
        self.assertNotIn("-U", command)
        self.assertNotIn("-P", command)

    def test_command_generation_with_windows_auth(self):
        """Test correct BCP command structure with Windows Authentication."""
        # Create connection config with Windows Auth
        windows_auth_config = SQLServerConnectionConfig(
            database="testdb",
            host="localhost",
            port=1433,
            use_windows_auth=True,
            auto_detect_driver=False,  # Don't validate driver existence in tests
        )

        # Create a new mock config that returns the Windows Auth config
        def mock_config_with_windows_auth(key):
            if key == "connections.source":
                return {self.engine: windows_auth_config}
            return None

        mock_program_config = Mock()
        mock_program_config.__getitem__ = Mock(side_effect=mock_config_with_windows_auth)

        data_source = BCPDataSource(
            engine=self.engine,
            statement="SELECT * FROM users",
            results_folder_path=self.temp_dir,
            logger=self.mock_logger,
            program_config=mock_program_config,
        )

        output_file = os.path.join(self.temp_dir, "result_001.csv")
        command = data_source._build_bcp_command(output_file)

        # Verify command structure
        self.assertEqual(command[0], "bcp")
        self.assertIn("SELECT * FROM users", command)
        self.assertIn("queryout", command)
        self.assertIn("-T", command)
        self.assertIn("-S", command)
        self.assertIn("-d", command)
        self.assertIn("testdb", command)


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