import boto3
from botocore.exceptions import ClientError
import uuid
import json
import zipfile
import os
import time

# A script to create an S3-triggered Lambda function for generating thumbnails using Boto3.

# --- Configuration ---
REGION = "us-east-1"
RANDOM_SUFFIX = uuid.uuid4().hex[:8]
SOURCE_BUCKET = f"my-source-images-boto3-{RANDOM_SUFFIX}"
DEST_BUCKET = f"my-destination-thumbnails-boto3-{RANDOM_SUFFIX}"
LAMBDA_FUNCTION_NAME = "S3ThumbnailGenerator-Boto3"
IAM_ROLE_NAME = "LambdaS3ThumbnailRole-Boto3"
POLICY_NAME = "LambdaS3ThumbnailPolicy-Boto3"
LAMBDA_CODE_FILE = "thumbnail_generator.py"
ZIP_FILE = "function.zip"

iam_client = boto3.client('iam', region_name=REGION)
s3_client = boto3.client('s3', region_name=REGION)
lambda_client = boto3.client('lambda', region_name=REGION)
sts_client = boto3.client('sts', region_name=REGION)

def create_iam_role():
    """Creates the IAM role and policy for the Lambda function."""
    print("Creating IAM Role and Policy...")
    
    trust_policy = {
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Principal": { "Service": "lambda.amazonaws.com" },
          "Action": "sts:AssumeRole"
        }
      ]
    }
    
    try:
        role_response = iam_client.create_role(
            RoleName=IAM_ROLE_NAME,
            AssumeRolePolicyDocument=json.dumps(trust_policy),
            Description="Role for S3 thumbnail generator Lambda"
        )
        role_arn = role_response['Role']['Arn']
        print(f"IAM Role created with ARN: {role_arn}")

        policy_document = {
            "Version": "2012-10-17",
            "Statement": [
                {"Effect": "Allow", "Action": ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"], "Resource": "arn:aws:logs:*:*:*"},
                {"Effect": "Allow", "Action": ["s3:GetObject"], "Resource": f"arn:aws:s3:::{SOURCE_BUCKET}/*"},
                {"Effect": "Allow", "Action": ["s3:PutObject"], "Resource": f"arn:aws:s3:::{DEST_BUCKET}/*"}
            ]
        }
        
        policy_response = iam_client.create_policy(
            PolicyName=POLICY_NAME,
            PolicyDocument=json.dumps(policy_document)
        )
        policy_arn = policy_response['Policy']['Arn']
        
        iam_client.attach_role_policy(RoleName=IAM_ROLE_NAME, PolicyArn=policy_arn)
        print(f"IAM Policy created and attached to the role.")
        
        # Wait for IAM propagation
        print("Waiting for IAM role to propagate...")
        time.sleep(10)
        
        return role_arn
    except ClientError as e:
        if e.response['Error']['Code'] == 'EntityAlreadyExists':
            print("IAM role already exists. Fetching ARN.")
            role_arn = iam_client.get_role(RoleName=IAM_ROLE_NAME)['Role']['Arn']
            return role_arn
        else:
            raise e

def create_s3_buckets():
    """Creates the source and destination S3 buckets."""
    print("Creating S3 buckets...")
    try:
        s3_client.create_bucket(Bucket=SOURCE_BUCKET, CreateBucketConfiguration={'LocationConstraint': REGION} if REGION != 'us-east-1' else {})
        s3_client.create_bucket(Bucket=DEST_BUCKET, CreateBucketConfiguration={'LocationConstraint': REGION} if REGION != 'us-east-1' else {})
        print(f"Source bucket created: {SOURCE_BUCKET}")
        print(f"Destination bucket created: {DEST_BUCKET}")
    except ClientError as e:
        print(f"Error creating buckets: {e}")
        raise

def create_lambda_function(role_arn):
    """Packages and creates the Lambda function."""
    print("Packaging and creating Lambda function...")
    try:
        # Package the Lambda code
        with zipfile.ZipFile(ZIP_FILE, 'w') as zf:
            zf.write(LAMBDA_CODE_FILE)
        
        with open(ZIP_FILE, 'rb') as f:
            zipped_code = f.read()

        response = lambda_client.create_function(
            FunctionName=LAMBDA_FUNCTION_NAME,
            Runtime='python3.9',
            Role=role_arn,
            Handler=f"{os.path.splitext(LAMBDA_CODE_FILE)[0]}.lambda_handler",
            Code={'ZipFile': zipped_code},
            Timeout=30,
            MemorySize=256,
            Environment={'Variables': {'DESTINATION_BUCKET': DEST_BUCKET}}
        )
        
        # Clean up local zip file
        os.remove(ZIP_FILE)
        
        lambda_arn = response['FunctionArn']
        print(f"Lambda function created with ARN: {lambda_arn}")
        return lambda_arn
    except ClientError as e:
        print(f"Error creating Lambda function: {e}")
        raise

def configure_s3_trigger(lambda_arn):
    """Configures the S3 bucket to trigger the Lambda function."""
    print("Configuring S3 trigger...")
    try:
        # Add permission for S3 to invoke Lambda
        lambda_client.add_permission(
            FunctionName=LAMBDA_FUNCTION_NAME,
            StatementId=f"S3InvokePermission-{RANDOM_SUFFIX}",
            Action='lambda:InvokeFunction',
            Principal='s3.amazonaws.com',
            SourceArn=f"arn:aws:s3:::{SOURCE_BUCKET}"
        )
        print("Lambda permission granted for S3 invocation.")

        # Configure bucket notification
        notification_config = {
            'LambdaFunctionConfigurations': [
                {
                    'LambdaFunctionArn': lambda_arn,
                    'Events': ['s3:ObjectCreated:*'],
                    'Filter': {
                        'Key': {
                            'FilterRules': [
                                {'Name': 'suffix', 'Value': '.jpg'},
                                {'Name': 'suffix', 'Value': '.png'}
                            ]
                        }
                    }
                }
            ]
        }
        
        s3_client.put_bucket_notification_configuration(
            Bucket=SOURCE_BUCKET,
            NotificationConfiguration=notification_config
        )
        print("S3 bucket notification configured.")
    except ClientError as e:
        print(f"Error configuring S3 trigger: {e}")
        raise

def main():
    """Main function to orchestrate the setup."""
    try:
        role_arn = create_iam_role()
        create_s3_buckets()
        lambda_arn = create_lambda_function(role_arn)
        configure_s3_trigger(lambda_arn)

        print("\n--- S3-triggered Lambda setup complete! ---")
        print(f"Upload a .jpg or .png file to the '{SOURCE_BUCKET}' bucket to trigger the function.")
        print(f"Check the '{DEST_BUCKET}' bucket for the generated thumbnail.")

    except Exception as e:
        print(f"Setup failed: {e}")

if __name__ == "__main__":
    main()