#!/usr/bin/env python3
"""
Test file for datacube S3 access with limited permissions
File test truy cập S3 qua datacube với quyền hạn chế
"""

import os
import sys
import boto3
from botocore.exceptions import ClientError, NoCredentialsError


def load_credentials(credential_file='train_files/crediential.txt'):
    """Load AWS credentials from file"""
    credentials = {}
    try:
        with open(credential_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith('export '):
                    line = line[7:]
                    if '=' in line:
                        key, value = line.split('=', 1)
                        value = value.strip('"')
                        credentials[key] = value
        
        print("✓ Credentials loaded successfully")
        print(f"  - AWS_ACCESS_KEY_ID: {credentials.get('AWS_ACCESS_KEY_ID', 'N/A')[:20]}...")
        return credentials
    except Exception as e:
        print(f"✗ Error loading credentials: {e}")
        return None


def set_aws_credentials(credentials):
    """Set AWS credentials as environment variables"""
    if not credentials:
        return False
    
    try:
        os.environ['AWS_ACCESS_KEY_ID'] = credentials.get('AWS_ACCESS_KEY_ID', '')
        os.environ['AWS_SECRET_ACCESS_KEY'] = credentials.get('AWS_SECRET_ACCESS_KEY', '')
        os.environ['AWS_SESSION_TOKEN'] = credentials.get('AWS_SESSION_TOKEN', '')
        print("✓ AWS credentials set in environment")
        return True
    except Exception as e:
        print(f"✗ Error setting credentials: {e}")
        return False


def test_s3_specific_bucket(bucket_name='deafrica-sentinel-2', region='ap-southeast-1', prefix=''):
    """
    Test access to a specific S3 bucket without requiring ListAllMyBuckets permission
    Kiểm tra truy cập bucket S3 cụ thể mà không cần quyền ListAllMyBuckets
    """
    try:
        print(f"\n=== Testing Specific Bucket Access ===")
        print(f"Bucket: {bucket_name}")
        print(f"Region: {region}")
        
        # Create S3 client
        s3_client = boto3.client('s3', region_name=region)
        
        # Try to access bucket with head_bucket (checks if bucket exists and we have access)
        try:
            s3_client.head_bucket(Bucket=bucket_name)
            print(f"✓ Successfully verified access to bucket: {bucket_name}")
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code', 'Unknown')
            if error_code == '404':
                print(f"✗ Bucket {bucket_name} does not exist or you don't have access")
                return None
            elif error_code == '403' or error_code == 'AccessDenied':
                print(f"✗ Access denied to bucket: {bucket_name}")
                return None
            else:
                print(f"✗ Error: {error_code}")
                return None
        
        # Try to list some objects
        print(f"\n✓ Attempting to list objects (max 5)...")
        response = s3_client.list_objects_v2(
            Bucket=bucket_name,
            MaxKeys=5,
            Prefix=prefix
        )
        
        if 'Contents' in response:
            print(f"✓ Found {len(response['Contents'])} objects:")
            for i, obj in enumerate(response['Contents'], 1):
                size_mb = obj['Size'] / (1024 * 1024)
                print(f"  {i}. {obj['Key']}")
                print(f"     Size: {size_mb:.2f} MB")
        else:
            print(f"✓ No objects found with prefix '{prefix}'")
        
        return s3_client
    
    except NoCredentialsError:
        print("\n✗ Error: No AWS credentials found")
        return None
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        print(f"\n✗ AWS Client Error: {e}")
        print(f"✗ Error Code: {error_code}")
        
        if error_code == 'ExpiredToken':
            print("✗ AWS session token has expired. Please refresh your credentials.")
        elif error_code == 'AccessDenied':
            print("✗ Access denied. You may not have permission for this bucket.")
        return None
    except Exception as e:
        print(f"\n✗ Unexpected error: {e}")
        return None


def test_datacube_s3_access():
    """
    Test S3 access using datacube pattern
    Kiểm tra truy cập S3 theo pattern của datacube
    """
    try:
        print("\n=== Testing Datacube S3 Access Pattern ===")
        
        # Import datacube S3 utilities
        try:
            from datacube.utils.rio import configure_s3_access
            print("✓ datacube library found")
            
            # Configure S3 access for rasterio/datacube
            aws_unsigned = False  # We have credentials
            region_name = 'ap-southeast-1'
            
            print(f"✓ Configuring S3 access for region: {region_name}")
            configure_s3_access(
                aws_unsigned=aws_unsigned,
                region_name=region_name,
                cloud_defaults=True
            )
            print("✓ S3 access configured for datacube/rasterio")
            
            return True
            
        except ImportError:
            print("✗ datacube library not found")
            print("  You can install it with: pip install datacube")
            return False
            
    except Exception as e:
        print(f"\n✗ Error configuring datacube S3 access: {e}")
        return False


def test_simple_s3_operations(bucket_name, s3_client):
    """
    Test basic S3 operations that work with limited permissions
    Kiểm tra các thao tác S3 cơ bản với quyền hạn chế
    """
    try:
        print("\n=== Testing Basic S3 Operations ===")
        
        # Test 1: Get bucket location
        try:
            response = s3_client.get_bucket_location(Bucket=bucket_name)
            location = response.get('LocationConstraint', 'us-east-1')
            print(f"✓ Bucket location: {location}")
        except ClientError as e:
            print(f"✗ Cannot get bucket location: {e.response.get('Error', {}).get('Code', 'Unknown')}")
        
        # Test 2: Check if we can read objects
        try:
            # Try to list with a common prefix
            response = s3_client.list_objects_v2(
                Bucket=bucket_name,
                MaxKeys=1,
                Delimiter='/'
            )
            
            if 'CommonPrefixes' in response:
                print(f"✓ Found {len(response['CommonPrefixes'])} top-level folders")
                for prefix in response['CommonPrefixes'][:3]:
                    print(f"  - {prefix['Prefix']}")
            
        except ClientError as e:
            print(f"✗ Cannot list objects: {e.response.get('Error', {}).get('Code', 'Unknown')}")
        
        return True
        
    except Exception as e:
        print(f"\n✗ Error in S3 operations: {e}")
        return False


def test_multiple_buckets():
    """
    Test access to multiple common S3 buckets
    Kiểm tra truy cập nhiều bucket S3 phổ biến
    """
    # Common buckets used in remote sensing / datacube projects
    test_buckets = [
        ('deafrica-sentinel-2', 'af-south-1', 'sentinel-s2-l2a-cogs'),
        ('sentinel-cogs', 'us-west-2', 'sentinel-s2-l2a'),
        ('usgs-landsat', 'us-west-2', 'collection02'),
        # EASI bucket might be private
        ('easi-asia-csiro', 'ap-southeast-1', ''),
    ]
    
    results = []
    
    for bucket_name, region, prefix in test_buckets:
        print(f"\n{'='*60}")
        print(f"Testing bucket: {bucket_name}")
        print(f"Region: {region}")
        
        s3_client = test_s3_specific_bucket(bucket_name, region, prefix)
        
        if s3_client:
            results.append((bucket_name, True))
            # If successful, try some operations
            test_simple_s3_operations(bucket_name, s3_client)
        else:
            results.append((bucket_name, False))
    
    return results


def main():
    """Main test function"""
    print("=" * 60)
    print("AWS S3 Datacube Access Test")
    print("Test Truy Cập S3 Datacube (với quyền hạn chế)")
    print("=" * 60)
    
    # Step 1: Load credentials
    print("\n[Step 1] Loading credentials from file...")
    credentials = load_credentials()
    
    if not credentials:
        print("\n✗ Test failed: Could not load credentials")
        return False
    
    # Step 2: Set credentials in environment
    print("\n[Step 2] Setting AWS credentials...")
    if not set_aws_credentials(credentials):
        print("\n✗ Test failed: Could not set credentials")
        return False
    
    # Step 3: Test datacube S3 access configuration
    print("\n[Step 3] Testing datacube S3 access configuration...")
    test_datacube_s3_access()
    
    # Step 4: Test access to multiple buckets
    print("\n[Step 4] Testing access to common S3 buckets...")
    results = test_multiple_buckets()
    
    # Summary
    print("\n" + "=" * 60)
    print("Test Summary / Tổng kết:")
    print("=" * 60)
    
    success_count = sum(1 for _, success in results if success)
    total_count = len(results)
    
    print(f"\nSuccessful connections: {success_count}/{total_count}")
    for bucket_name, success in results:
        status = "✓" if success else "✗"
        print(f"{status} {bucket_name}")
    
    print("\n" + "=" * 60)
    if success_count > 0:
        print("✓ Test completed! You have access to some buckets.")
        print("✓ Test hoàn thành! Bạn có quyền truy cập một số bucket.")
    else:
        print("⚠ No buckets accessible with current credentials.")
        print("⚠ Không có bucket nào truy cập được với credentials hiện tại.")
    print("=" * 60)
    
    return success_count > 0


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