import pytest
from unittest.mock import Mock, MagicMock, patch
import pandas as pd
from io import BytesIO
import zipfile
from pydataframer_databricks import FileType, DatasetType
from pydataframer_databricks import FileType


class TestDatabricksConnector:
    """Test suite for DatabricksConnector class"""
    
    @pytest.fixture
    def mock_dbutils(self):
        """Create a mock dbutils object"""
        dbutils = Mock()
        dbutils.secrets.get = Mock(side_effect=lambda scope, key: {
            ("dataframer", "DATABRICKS_SERVER_HOSTNAME"): "test.databricks.com",
            ("dataframer", "DATABRICKS_HTTP_PATH"): "/sql/1.0/warehouses/abc123",
            ("dataframer", "DATABRICKS_CLIENT_ID"): "test-client-id",
            ("dataframer", "DATABRICKS_CLIENT_SECRET"): "test-client-secret",
        }.get((scope, key), "default-value"))
        return dbutils
    
    @pytest.fixture
    def connector(self, mock_dbutils):
        """Create a DatabricksConnector instance with mocked dbutils"""
        from pydataframer_databricks import DatabricksConnector
        return DatabricksConnector(mock_dbutils)
    
    def test_init(self, mock_dbutils):
        """Test connector initialization"""
        from pydataframer_databricks import DatabricksConnector
        connector = DatabricksConnector(mock_dbutils)
        assert connector.dbutils == mock_dbutils
    
    def test_get_connection(self, connector, mock_dbutils):
        """Test get_connection establishes connection with correct parameters"""
        with patch('databricks.sql.connect') as mock_sql_connect:
            mock_connection = Mock()
            mock_sql_connect.return_value = mock_connection
            
            with patch('databricks.sdk.core.oauth_service_principal') as mock_oauth:
                result = connector.get_connection()
                
                mock_dbutils.secrets.get.assert_called()
                
                mock_sql_connect.assert_called_once()
                call_kwargs = mock_sql_connect.call_args.kwargs
                assert call_kwargs['server_hostname'] == "test.databricks.com"
                assert call_kwargs['http_path'] == "/sql/1.0/warehouses/abc123"
                assert call_kwargs['user_agent_entry'] == "dataframer_user_agent"
                
                assert result == mock_connection
    
    @patch('pydataframer_databricks.connectors.DatabricksConnector.get_connection')
    def test_fetch_sample_data_success(self, mock_get_connection, connector):
        """Test fetch_sample_data successfully fetches and returns DataFrame"""
        mock_cursor = Mock()
        mock_cursor.description = [("id",), ("name",), ("value",)]
        mock_cursor.fetchall.return_value = [
            (1, "test1", 100),
            (2, "test2", 200),
        ]
        
        mock_connection = MagicMock()
        mock_connection.__enter__.return_value = mock_connection
        mock_connection.cursor.return_value.__enter__.return_value = mock_cursor
        
        mock_get_connection.return_value = mock_connection
        
        result = connector.fetch_sample_data(
            num_items_to_select=25,
            table_name="test_catalog.test_schema.test_table"
        )
        
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 2
        assert list(result.columns) == ["id", "name", "value"]
        mock_cursor.execute.assert_called_once()
    
    @patch('pydataframer_databricks.connectors.DatabricksConnector.get_connection')
    def test_fetch_sample_data_query_failure(self, mock_get_connection, connector):
        """Test fetch_sample_data handles query execution errors"""
        mock_cursor = Mock()
        mock_cursor.execute.side_effect = Exception("Table not found")
        
        mock_connection = MagicMock()
        mock_connection.__enter__.return_value = mock_connection
        mock_connection.cursor.return_value.__enter__.return_value = mock_cursor
        
        mock_get_connection.return_value = mock_connection
        
        with pytest.raises(RuntimeError) as exc_info:
            connector.fetch_sample_data(
                num_items_to_select=10,
                table_name="nonexistent.table"
            )
        
        assert "Failed to fetch data from table" in str(exc_info.value)
    
    def test_load_generated_data_single_file_csv_success(self, connector):
        """Test load_generated_data with SINGLE_FILE CSV dataset"""
        csv_content = b"id,name,value\n1,test1,100\n2,test2,200\n"
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as z:
            z.writestr("generated_samples.csv", csv_content)
        zip_buffer.seek(0)
        
        mock_cursor = Mock()
        mock_connection = MagicMock()
        mock_connection.__enter__.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        
        with patch.object(connector, 'get_connection', return_value=mock_connection):
            connector.load_generated_data(
                table_name="test_catalog.test_schema.test_table",
                downloaded_zip=zip_buffer,
                dataset_type=DatasetType.SINGLE_FILE,
                file_type=FileType.CSV
            )
        
        assert mock_cursor.execute.call_count == 1
        create_table_call = mock_cursor.execute.call_args[0][0]
        assert "CREATE OR REPLACE TABLE" in create_table_call
        assert "test_catalog.test_schema.test_table" in create_table_call
        
        assert mock_cursor.executemany.call_count == 1
        assert mock_cursor.close.call_count == 1
    
    def test_load_generated_data_invalid_zip(self, connector):
        """Test load_generated_data handles invalid ZIP files"""
        invalid_zip = BytesIO(b"not a zip file")
        
        with pytest.raises(ValueError) as exc_info:
            connector.load_generated_data(
                table_name="test_table",
                downloaded_zip=invalid_zip,
                dataset_type=DatasetType.SINGLE_FILE,
                file_type=FileType.CSV
            )
        
        assert "Invalid or corrupted ZIP file" in str(exc_info.value)
    
    def test_load_generated_data_no_csv_file(self, connector):
        """Test load_generated_data when ZIP has no CSV file"""
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as z:
            z.writestr("data.txt", b"some text")
        zip_buffer.seek(0)
        
        with pytest.raises(ValueError) as exc_info:
            connector.load_generated_data(
                table_name="test_table",
                downloaded_zip=zip_buffer,
                dataset_type=DatasetType.SINGLE_FILE,
                file_type=FileType.CSV
            )
        
        assert "Expected exactly one .csv file in ZIP" in str(exc_info.value)
    
    def test_load_generated_data_multiple_csv_files(self, connector):
        """Test load_generated_data when ZIP has multiple CSV files"""
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as z:
            z.writestr("data1.csv", b"id,name\n1,test1")
            z.writestr("data2.csv", b"id,name\n2,test2")
        zip_buffer.seek(0)
        
        with pytest.raises(ValueError) as exc_info:
            connector.load_generated_data(
                table_name="test_table",
                downloaded_zip=zip_buffer,
                dataset_type=DatasetType.SINGLE_FILE,
                file_type=FileType.CSV
            )
        
        assert "Expected exactly one .csv file in ZIP" in str(exc_info.value)
    
    def test_load_generated_data_create_table_failure(self, connector):
        """Test load_generated_data handles CREATE TABLE errors"""
        csv_content = b"id,name\n1,test1"
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as z:
            z.writestr("data.csv", csv_content)
        zip_buffer.seek(0)
        
        mock_cursor = Mock()
        mock_cursor.execute.side_effect = Exception("Permission denied")
        mock_connection = MagicMock()
        mock_connection.__enter__.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        
        with patch.object(connector, 'get_connection', return_value=mock_connection):
            with pytest.raises(RuntimeError) as exc_info:
                connector.load_generated_data(
                    table_name="test_table",
                    downloaded_zip=zip_buffer,
                    dataset_type=DatasetType.SINGLE_FILE,
                    file_type=FileType.CSV
                )
        
        assert "Failed to create table" in str(exc_info.value)
        assert mock_cursor.close.call_count == 1
    
    def test_load_generated_data_insert_failure(self, connector):
        """Test load_generated_data handles INSERT errors"""
        csv_content = b"id,name\n1,test1"
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w') as z:
            z.writestr("data.csv", csv_content)
        zip_buffer.seek(0)
        
        mock_cursor = Mock()
        mock_cursor.execute.return_value = None
        mock_cursor.executemany.side_effect = Exception("Constraint violation")
        mock_connection = MagicMock()
        mock_connection.__enter__.return_value = mock_connection
        mock_connection.cursor.return_value = mock_cursor
        
        with patch.object(connector, 'get_connection', return_value=mock_connection):
            with pytest.raises(RuntimeError) as exc_info:
                connector.load_generated_data(
                    table_name="test_table",
                    downloaded_zip=zip_buffer,
                    dataset_type=DatasetType.SINGLE_FILE,
                    file_type=FileType.CSV
                )
        
        assert "Failed to insert data" in str(exc_info.value)
        assert mock_cursor.close.call_count == 1
