import boto3
import re
from collections import defaultdict
import time

# === User Configuration ===
BUCKET = 'o11y-prod-shared-us-west-2-staging'
PREFIX_ROOT = '1/'
START_TIMESTAMP = 1785196800
END_TIMESTAMP = 1785283200

s3 = boto3.client('s3')
timestamp_pattern = re.compile(r'(\d{10})-')
stats = defaultdict(lambda: {'count': 0, 'total_size': 0})


def list_prefixes(bucket, prefix):
    """List all directories (sub-prefixes) under the next level of prefix"""
    paginator = s3.get_paginator('list_objects_v2')
    #print("list prefix: " + prefix)
    pages = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter='/')
    for page in pages:
        for cp in page.get('CommonPrefixes', []):
            #print("get page: " + cp['Prefix'])
            yield cp['Prefix']


def write_stats_to_file(prefix, count, size_mb):
    """Write statistics to file"""
    filename = f"scan_stats_{BUCKET}_1_{START_TIMESTAMP}_{END_TIMESTAMP}.txt"
    with open(filename, 'a') as f:
        line = f"{prefix},{count},{size_mb}\n"
        f.write(line)
    #print(f"Stats written to {filename}")

def scan_logs(bucket, prefix):
    """List and count objects within the specified time range under logs/{service}/
    
    Use StartAfter parameter to efficiently list objects within the specified prefix and time range
    """
    paginator = s3.get_paginator('list_objects_v2')
    
    # Construct start and end prefixes
    start_prefix = f"{prefix}{START_TIMESTAMP}"
    end_prefix = f"{prefix}{END_TIMESTAMP + 1}"  # Add 1 to ensure end timestamp is included
    
    # Use paginator to list all objects
    pages = paginator.paginate(
        Bucket=bucket,
        Prefix=prefix,
        StartAfter=start_prefix,
        PaginationConfig={
            'PageSize': 1000,
            'StartingToken': None
        }
    )
    count = 0
    try:
        for page in pages:
            if 'Contents' not in page:
                continue
                
            for obj in page['Contents']:
                key = obj['Key']
                # Stop processing if object key exceeds end timestamp
                if key >= end_prefix:
                    size_mb = stats[prefix]['total_size'] / (1024 * 1024)  # Convert to MB
                    print(f"scan logs done: {prefix}{START_TIMESTAMP} count: {count} size: {size_mb:.2f} MB")
                    write_stats_to_file(prefix, count, f"{size_mb:.2f}")
                    return
                    
                size = obj['Size']
                count += 1
                if count % 10000 == 0:
                    print(f"scan logs: {count} key: {key}")
                stats[prefix]['count'] += 1
                stats[prefix]['total_size'] += size

        size_mb = stats[prefix]['total_size'] / (1024 * 1024)  # Convert to MB
        print(f"scan logs done: {prefix}{START_TIMESTAMP} count: {count} size: {size_mb:.2f} MB")
        write_stats_to_file(prefix, count, f"{size_mb:.2f}")
        
    except Exception as e:
        size_mb = stats[prefix]['total_size'] / (1024 * 1024)  # Convert to MB
        print(f"Error during scan: {str(e)}")
        print(f"Partial results - count: {count}, size: {size_mb:.2f} MB")
        write_stats_to_file(f"ERROR_{prefix}", count, f"{size_mb:.2f}")
        raise


# === Recursively 5 levels of directories (including logs/) ===
for p1 in list_prefixes(BUCKET, PREFIX_ROOT):          # 0/xxx/
    for p2 in list_prefixes(BUCKET, p1):               # 0/xxx/yyy/
        for p3 in list_prefixes(BUCKET, p2):           # ...
            for p4 in list_prefixes(BUCKET, p3):
                if p4.endswith('/logs/'):
                    for p5 in list_prefixes(BUCKET, p4):   # To logs/
                        #print("get p5: " + p5)
                        scan_logs(BUCKET, p5)
                else:
                    scan_logs(BUCKET, p4)

# === Print statistics ===
print(f"{'Prefix':80} | Files | Total Size (MB)")
print("-" * 100)
for prefix, data in sorted(stats.items()):
    size_mb = data['total_size'] / 1024 / 1024
    print(f"{prefix:80} | {data['count']:5d} | {size_mb:12.2f}")