import unittest
import os
from unittest.mock import patch, MagicMock

# Only import if available
try:
    from deltalake import DeltaTable
    from o11ydiagnosis.subtasks.slowlogconverter.log import log as logger
    DELTALAKE_AVAILABLE = True
except ImportError:
    DELTALAKE_AVAILABLE = False
    logger = MagicMock()

def optimize_table(region: str, s3_path: str, with_date: bool):
    storage_options = {
        'AWS_REGION': region,
    }

    try:
        logger.info(f"Attempting to optimize (compact) slowlogs table at {s3_path}")

        delta_table = DeltaTable(s3_path, storage_options=storage_options)
        logger.info(f"Starting compaction {s3_path}")
        logger.info(f"Table version: {delta_table.version()}")
        if with_date:
            optimize_metrics = delta_table.optimize.compact(partition_filters=[("date", ">=", "20250501"), ("date", "<", "20250630")])
        else:
            optimize_metrics = delta_table.optimize.compact()

        logger.info(f"Slow Query table compacted. Metrics: {optimize_metrics}")

        logger.info(f"Starting vacuum {s3_path}")

        delta_table = DeltaTable(s3_path, storage_options=storage_options)
        logger.info(f"Table version: {delta_table.version()}")
        vacuum_metrics = delta_table.vacuum(retention_hours=2, enforce_retention_duration=False, dry_run=False)
        logger.info(f"Slow Query table vacuum. Metrics: {vacuum_metrics}")

    except Exception as e:
        logger.error(f"Error optimizing slow query table at {s3_path}: {e}")

class TestOptimizeTable(unittest.TestCase):
    @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake not available")
    @unittest.skipUnless(os.environ.get("AWS_ACCESS_KEY_ID"), "AWS credentials not set")
    def test_optimize_table_with_date(self):
        """Test optimize_table with date partition filter"""
        region = "us-east-1"
        s3_path = "s3://test-bucket/test-path/slowlogs"
        optimize_table(region, s3_path, True)
        # If no exception, test passes
    
    @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake not available")
    @unittest.skipUnless(os.environ.get("AWS_ACCESS_KEY_ID"), "AWS credentials not set")
    def test_optimize_table_without_date(self):
        """Test optimize_table without date partition filter"""
        region = "us-east-1"
        s3_path = "s3://test-bucket/test-path/plans"
        optimize_table(region, s3_path, False)
        # If no exception, test passes
    
    @patch('o11ydiagnosis.subtasks.slowlogconverter.test_query_v2.DeltaTable')
    def test_optimize_table_mock(self, mock_delta_table):
        """Test optimize_table with mocked DeltaTable"""
        mock_table = MagicMock()
        mock_table.version.return_value = 1
        mock_table.optimize.compact.return_value = {"files_added": 0, "files_removed": 0}
        mock_table.vacuum.return_value = {"files_deleted": 0}
        mock_delta_table.return_value = mock_table
        
        region = "us-east-1"
        s3_path = "s3://test-bucket/test-path/slowlogs"
        optimize_table(region, s3_path, True)
        
        # Verify DeltaTable was called
        mock_delta_table.assert_called()
        mock_table.optimize.compact.assert_called()
        mock_table.vacuum.assert_called()


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