#!/usr/bin/env python3
"""
Test S3 list all objects
Test liệt kê tất cả objects trong S3 bucket
"""

import boto3
from botocore.exceptions import ClientError
from cognito_auth import CognitoAuthenticator


def list_all_s3_objects(bucket_name, region='us-west-2', prefix='', max_keys=0):
    """
    List all objects in S3 bucket
    
    Args:
        bucket_name: Tên bucket
        region: AWS region  
        prefix: Prefix để filter
        max_keys: Số lượng max (0 = tất cả)
    """
    print("=" * 70)
    print(f"Listing S3 Objects")
    print(f"Bucket: {bucket_name}")
    print(f"Region: {region}")
    if prefix:
        print(f"Prefix: {prefix}")
    print(f"Max: {max_keys if max_keys > 0 else 'ALL'}")
    print("=" * 70)
    
    # Initialize auth
    auth = CognitoAuthenticator(region='ap-southeast-1')
    
    # Load credentials
    if not auth.load_tokens_from_file('train_files/crediential.txt'):
        print("✗ Failed to load credentials")
        return []
    
    auth.get_credentials_from_cognito()
    
    # Create S3 client
    s3_client = boto3.client(
        's3',
        region_name=region,
        aws_access_key_id=auth.aws_credentials['AccessKeyId'],
        aws_secret_access_key=auth.aws_credentials['SecretAccessKey'],
        aws_session_token=auth.aws_credentials['SessionToken']
    )
    
    # List objects with pagination
    all_objects = []
    continuation_token = None
    page = 0
    
    try:
        while True:
            page += 1
            print(f"\n📄 Page {page}...")
            
            # Prepare parameters
            list_params = {
                'Bucket': bucket_name,
                'MaxKeys': 1000 if max_keys == 0 else min(max_keys - len(all_objects), 1000)
            }
            
            if prefix:
                list_params['Prefix'] = prefix
                
            if continuation_token:
                list_params['ContinuationToken'] = continuation_token
            
            # List objects
            response = s3_client.list_objects_v2(**list_params)
            
            if 'Contents' in response:
                page_objects = response['Contents']
                all_objects.extend(page_objects)
                
                # Show some samples from this page
                print(f"  Found {len(page_objects)} objects on this page")
                for obj in page_objects[:3]:
                    size_mb = obj['Size'] / (1024 * 1024)
                    print(f"    - {obj['Key']} ({size_mb:.2f} MB)")
                
                if len(page_objects) > 3:
                    print(f"    ... and {len(page_objects) - 3} more")
            
            # Check if should continue
            if max_keys > 0 and len(all_objects) >= max_keys:
                print(f"\n✓ Reached max_keys limit: {max_keys}")
                break
                
            if not response.get('IsTruncated', False):
                print(f"\n✓ Reached end of list")
                break
                
            continuation_token = response.get('NextContinuationToken')
        
        # Summary
        print("\n" + "=" * 70)
        print("Summary / Tổng Kết")
        print("=" * 70)
        
        if all_objects:
            total_size = sum(obj['Size'] for obj in all_objects)
            total_size_gb = total_size / (1024 * 1024 * 1024)
            
            print(f"\n✓ Total objects: {len(all_objects)}")
            print(f"✓ Total size: {total_size_gb:.2f} GB")
            print(f"✓ Pages fetched: {page}")
            
            # Show first and last
            print(f"\nFirst 5 objects:")
            for i, obj in enumerate(all_objects[:5], 1):
                size_mb = obj['Size'] / (1024 * 1024)
                print(f"  {i}. {obj['Key']}")
                print(f"     Size: {size_mb:.2f} MB | Modified: {obj['LastModified']}")
            
            if len(all_objects) > 10:
                print(f"\nLast 5 objects:")
                for i, obj in enumerate(all_objects[-5:], len(all_objects)-4):
                    size_mb = obj['Size'] / (1024 * 1024)
                    print(f"  {i}. {obj['Key']}")
                    print(f"     Size: {size_mb:.2f} MB | Modified: {obj['LastModified']}")
        else:
            print("\n⚠ No objects found")
        
        print("=" * 70)
        
        return all_objects
        
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        print(f"\n✗ S3 Error: {error_code}")
        print(f"✗ Message: {e}")
        return []
    except Exception as e:
        print(f"\n✗ Error: {e}")
        import traceback
        traceback.print_exc()
        return []


def main():
    """Test listing objects"""
    
    print("\n" + "╔" + "═" * 68 + "╗")
    print("║" + " " * 20 + "S3 OBJECT LISTING TEST" + " " * 26 + "║")
    print("╚" + "═" * 68 + "╝\n")
    
    # Test cases
    test_cases = [
        {
            'bucket': 'sentinel-cogs',
            'region': 'us-west-2',
            'prefix': 'sentinel-s2-l2a-cogs/54/S/VE/2020/1/',
            'max_keys': 100,
            'description': 'Sentinel-2 Vietnam 2020 January data'
        },
        {
            'bucket': 'sentinel-cogs',
            'region': 'us-west-2',
            'prefix': 'sentinel-s2-l2a-cogs/',
            'max_keys': 50,
            'description': 'Sentinel-2 global data (sample)'
        },
    ]
    
    for i, test in enumerate(test_cases, 1):
        print(f"\n{'='*70}")
        print(f"TEST CASE {i}: {test['description']}")
        print(f"{'='*70}\n")
        
        objects = list_all_s3_objects(
            bucket_name=test['bucket'],
            region=test['region'],
            prefix=test['prefix'],
            max_keys=test['max_keys']
        )
        
        input(f"\n⏸  Press Enter to continue to next test...")
    
    print("\n✓ All tests completed!")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n✗ Interrupted by user")
    except Exception as e:
        print(f"\n✗ Fatal error: {e}")
        import traceback
        traceback.print_exc()