#!/usr/bin/env python3
"""
Test S3 access using Cognito tokens
Test truy cập S3 sử dụng Cognito tokens
"""

import sys
from cognito_auth import CognitoAuthenticator
from datacube.utils.rio import configure_s3_access


def test_cognito_to_s3():
    """
    Complete test: Cognito tokens → AWS credentials → S3 access
    """
    print("=" * 70)
    print("Test S3 Access Using Cognito Tokens")
    print("Test Truy Cập S3 Sử Dụng Cognito Tokens")
    print("=" * 70)
    
    # Initialize
    auth = CognitoAuthenticator(region='ap-southeast-1')
    
    # Load Cognito tokens
    print("\n[1/5] Loading Cognito tokens from file...")
    if not auth.load_tokens_from_file('train_files/crediential.txt'):
        print("✗ Failed to load tokens")
        return False
    
    # Show token info
    print("\n[2/5] Displaying token information...")
    decoded_id, decoded_access = auth.print_token_info()
    
    if not decoded_id:
        print("✗ Invalid tokens")
        return False
    
    # Get AWS credentials from Cognito
    print("\n[3/5] Getting AWS credentials...")
    print("ℹ Note: Since we don't have Identity Pool ID, using existing credentials")
    
    # Use existing credentials from file (already exchanged by EASI)
    if not auth.get_credentials_from_cognito():
        print("✗ Failed to get AWS credentials")
        return False
    
    # Set credentials in environment
    print("\n[4/5] Configuring environment...")
    if not auth.set_environment_credentials():
        print("✗ Failed to set environment")
        return False
    
    # Configure datacube S3 access
    print("\n[5/5] Configuring datacube S3 access...")
    try:
        configure_s3_access(
            aws_unsigned=False,
            region_name='us-west-2',
            cloud_defaults=True
        )
        print("✓ Datacube S3 access configured")
    except Exception as e:
        print(f"⚠ Warning: Could not configure datacube: {e}")
    
    # Test multiple buckets
    print("\n" + "=" * 70)
    print("Testing S3 Bucket Access")
    print("=" * 70)
    
    test_buckets = [
        ('sentinel-cogs', 'us-west-2', 'sentinel-s2-l2a-cogs/'),
        ('sentinel-s2-l2a', 'eu-central-1', ''),
    ]
    
    results = []
    for bucket_name, region, prefix in test_buckets:
        print(f"\nTesting: {bucket_name} ({region})")
        print("-" * 50)
        # List up to 50 objects (có thể thay đổi số này)
        success = auth.test_s3_access(bucket_name, region, max_keys=50, prefix=prefix)
        results.append((bucket_name, success))
    
    # Summary
    print("\n" + "=" * 70)
    print("Summary / Tổng Kết")
    print("=" * 70)
    
    print("\nAuthentication Flow:")
    print("  Cognito Tokens → ✓")
    print("  AWS Credentials → ✓")
    print("  Environment Setup → ✓")
    
    print(f"\nS3 Access Results:")
    success_count = sum(1 for _, success in results if success)
    for bucket, success in results:
        status = "✓" if success else "✗"
        print(f"  {status} {bucket}")
    
    print(f"\nTotal: {success_count}/{len(results)} buckets accessible")
    
    if success_count > 0:
        print("\n✓ SUCCESS: Cognito authentication working!")
        print("✓ THÀNH CÔNG: Xác thực Cognito hoạt động!")
    else:
        print("\n⚠ WARNING: Could not access any S3 buckets")
        print("⚠ CẢNH BÁO: Không thể truy cập bucket S3 nào")
    
    print("=" * 70)
    
    return success_count > 0


def main():
    try:
        success = test_cognito_to_s3()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n✗ Test interrupted")
        sys.exit(1)
    except Exception as e:
        print(f"\n✗ Fatal error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()