import boto3
import json
import uuid
import argparse
from constants import AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY

def test_s3_credentials(read_only=False):
    """Test S3 credentials by reading and writing to the sync token cache"""
    print("Testing AWS S3 credentials for sync token cache...")
    
    # Initialize S3 client with credentials
    s3_client = boto3.client(
        "s3", 
        aws_access_key_id=AWS_ACCESS_KEY_ID,
        aws_secret_access_key=AWS_SECRET_ACCESS_KEY
    )
    
    bucket = "suno-alexa-sync-token-cache"
    key = "user_id_sync_token.json"
    test_user_id = f"test-user-{uuid.uuid4()}"
    test_token = str(uuid.uuid4())
    
    # Step 1: Try to read the existing file (if any)
    print("\n--- Testing Read Access ---")
    try:
        response = s3_client.get_object(
            Bucket=bucket,
            Key=key
        )
        existing_data = json.loads(response["Body"].read().decode("utf-8"))
        print(f"✅ Successfully read cache file")
        print(f"Current cache contains {len(existing_data)} entries")
        
        # Pretty print the current contents
        print("\n--- Current Cache Contents ---")
        print(json.dumps(existing_data, indent=2, sort_keys=True))
        
    except s3_client.exceptions.NoSuchKey:
        print("Cache file doesn't exist yet - will create a new one")
        existing_data = {}
    except Exception as e:
        print(f"❌ Error reading cache file: {str(e)}")
        existing_data = {}
    
    # Skip write testing if read_only mode is enabled
    if read_only:
        print("\n--- Write Testing Skipped (Read-Only Mode) ---")
        return
    
    # Step 2: Add test data and write back
    print("\n--- Testing Write Access ---")
    # Make a backup of the original data
    original_data = existing_data.copy()
    
    try:
        # Add test entry
        existing_data[test_user_id] = test_token
        
        # Write back to S3
        s3_client.put_object(
            Bucket=bucket,
            Key=key,
            Body=json.dumps(existing_data),
            ContentType="application/json"
        )
        print(f"✅ Successfully wrote test entry to cache")
        
        # Read it back to verify
        response = s3_client.get_object(
            Bucket=bucket,
            Key=key
        )
        verification_data = json.loads(response["Body"].read().decode("utf-8"))
        retrieved_token = verification_data.get(test_user_id)
        
        if retrieved_token == test_token:
            print(f"✅ Successfully verified test entry in cache")
        else:
            print(f"❌ Verification failed - token mismatch")
    
    except Exception as e:
        print(f"❌ Error during write/verify test: {str(e)}")
    
    finally:
        # Always clean up test entry and restore original data
        print("\n--- Cleaning Up Test Data ---")
        try:
            s3_client.put_object(
                Bucket=bucket,
                Key=key,
                Body=json.dumps(original_data),
                ContentType="application/json"
            )
            print(f"✅ Successfully restored original cache data")
        except Exception as e:
            print(f"❌ Error during cleanup: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test S3 credentials and sync token cache")
    parser.add_argument("--read-only", action="store_true", help="Only read the cache, don't write any test data")
    args = parser.parse_args()
    
    test_s3_credentials(read_only=args.read_only)