import boto3
import gzip
import os
import tempfile
from botocore.exceptions import ClientError

# Configuration
BUCKET = 'o11y-prod-shared-ap-northeast-1'  # Update this to your bucket name
PREFIX_ROOT = 'diagnosis/data/'

# Initialize S3 client
s3 = boto3.client('s3')

def list_prefixes(bucket, prefix, delimiter='/'):
    """List all prefixes in the given directory."""
    paginator = s3.get_paginator('list_objects_v2')
    result = paginator.paginate(
        Bucket=bucket,
        Prefix=prefix,
        Delimiter=delimiter
    )
    for page in result:
        if 'CommonPrefixes' in page:
            for obj in page['CommonPrefixes']:
                yield obj['Prefix']

def list_objects(bucket, prefix):
    """List all objects with the given prefix."""
    paginator = s3.get_paginator('list_objects_v2')
    result = paginator.paginate(
        Bucket=bucket,
        Prefix=prefix,
        PaginationConfig={
            'PageSize': 1000,
            'StartingToken': None
        }
    )
    for page in result:
        if 'Contents' in page:
            for obj in page['Contents']:
                yield obj

def is_gzipped(data):
    """Check if the given data is gzipped by checking the magic number."""
    return len(data) > 1 and data[0] == 0x1F and data[1] == 0x8B

def compress_and_upload(s3_key):
    """Download, compress and re-upload the file if it's not already compressed.
    
    Returns:
        bool: True if file was processed (compressed or already compressed), False if skipped or error occurred.
    """
    # Skip if already has .gz extension
    if s3_key.endswith('.gz'):
        print(f"Skipping already compressed file: {s3_key}")
        return True

    try:
        print(f"Processing: {s3_key}")
        
        # Create a temporary file in current directory and download the content
        with tempfile.NamedTemporaryFile(delete=False, dir='.') as tmp_file:
            tmp_file_path = tmp_file.name
            # Download the file
            s3.download_fileobj(BUCKET, s3_key, tmp_file)
            
        # Check if the downloaded file is already gzipped
        with open(tmp_file_path, 'rb') as f:
            first_two_bytes = f.read(2)
            if is_gzipped(first_two_bytes):
                print(f"Skipping already gzipped file (by magic number): {s3_key}")
                os.unlink(tmp_file_path)  # Clean up the temp file
                return True

        # Compress the file
        compressed_path = f"{tmp_file_path}.gz"
        with open(tmp_file_path, 'rb') as f_in:
            with gzip.open(compressed_path, 'wb') as f_out:
                f_out.writelines(f_in)

        # Upload the compressed file with the original key (overwrite)
        s3.upload_file(compressed_path, BUCKET, s3_key)
        print(f"Uploaded compressed file: {s3_key}")

        # Clean up
        os.unlink(tmp_file_path)
        os.unlink(compressed_path)
        
        return True

    except ClientError as e:
        print(f"Error processing {s3_key}: {e}")
        return False
    except Exception as e:
        print(f"Unexpected error with {s3_key}: {e}")
        return False
    finally:
        # Clean up in case of any errors
        if 'tmp_file_path' in locals() and os.path.exists(tmp_file_path):
            os.unlink(tmp_file_path)
        if 'compressed_path' in locals() and os.path.exists(compressed_path):
            os.unlink(compressed_path)

def process_directory(prefix):
    """Process all files in the given directory."""
    processed_count = 0
    
    for obj in list_objects(BUCKET, prefix):
        if not obj['Key'].endswith('/'):  # Skip directory markers
            if compress_and_upload(obj['Key']):
                processed_count += 1
            else:
                print(f"processed to compress {obj['Key']}, done")
                break
    return processed_count

def main():
    print("Starting compression process...")
    total_processed = 0
    
    # Process all directories under the root prefix
    # for p1 in list_prefixes(BUCKET, PREFIX_ROOT):          # 0/xxx/               # 0/xxx/yyy/
    #     if p1 < 'diagnosis/data/10408003622347968695/' or \
    #     p1.startswith('diagnosis/data/1379661944644514084/') or \
    #     p1.startswith('diagnosis/data/10947714666900611238/') or \
    #     p1.startswith('diagnosis/data/10665115838626244426/') or \
    #     p1.startswith('diagnosis/data/10795788448421792248/') or \
    #     p1.startswith('diagnosis/data/1379661944646413605/') or \
    #     p1.startswith('diagnosis/data/seed/') or \
    #     p1.startswith('diagnosis/data/o11y/') or \
    #     p1.startswith('diagnosis/data/1379661944646414964/'):
    #         print(f"skip {p1}")
    #         continue  
    p1 = 'diagnosis/data/o11y/'
    for p3 in list_prefixes(BUCKET, p1+"merged-logs/"):
        total_processed += process_directory(p3)
    
    print(f"\nCompression process completed. Total files processed: {total_processed}")

if __name__ == "__main__":
    main()