#!/usr/bin/env python3
import argparse
import logging
import os
import sys
from datetime import datetime, date, timedelta
from deltalake import DeltaTable

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def optimize_delta_table(s3_path, region):
    """
    Optimize and vacuum a Delta table in S3.
    
    Args:
        s3_path (str): Full S3 path to the Delta table (e.g., 's3://bucket/prefix/')
        region (str): AWS region where the S3 bucket is located
        partition_date (str, optional): Date in 'YYYYMMDD' format for partition optimization
    """
    storage_options = {'AWS_REGION': region}
    
    try:
        logger.info(f"Connecting to Delta table at {s3_path}")
        delta_table = DeltaTable(s3_path, storage_options=storage_options)
        logger.info(f"Table version: {delta_table.version()}")
        
        logger.info(f"Optimizing full table: {s3_path}")
        optimize_metrics = delta_table.optimize.compact(target_size=512 * 1024 * 1024)

        logger.info(f"Table compacted. Metrics: {optimize_metrics}")

        logger.info(f"Starting vacuum for {s3_path}")
        delta_table = DeltaTable(s3_path, storage_options=storage_options)
        logger.info(f"Table version: {delta_table.version()}")
        delta_table.alter.set_table_properties({"delta.logRetentionDuration": "interval 2 hours"})
        delta_table.create_checkpoint()
        vacuum_metrics = delta_table.vacuum(
            retention_hours=2, 
            enforce_retention_duration=False, 
            dry_run=False
        )
        delta_table.cleanup_metadata()
        logger.info(f"Table vacuum completed. Metrics: {vacuum_metrics}")
        return 0
        
    except Exception as e:
        logger.error(f"Error optimizing table {s3_path}: {str(e)}", exc_info=True)
        return 1

def main():
    parser = argparse.ArgumentParser(description='Optimize and vacuum a Delta table in S3.')
    parser.add_argument('--s3-path', required=True, help='S3 path to the Delta table (e.g., s3://bucket/prefix/)')
    parser.add_argument('--region', required=True, help='AWS region where the S3 bucket is located')
    
    args = parser.parse_args()
    
    # Ensure s3_path ends with a slash
    s3_path = args.s3_path.rstrip('/') + '/'
    
    return optimize_delta_table(s3_path, args.region)

if __name__ == "__main__":
    sys.exit(main())