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

# A script to create a simple API Gateway REST API with a Lambda proxy integration using Boto3.

# --- Configuration ---
REGION = "us-east-1"
API_NAME = "MyLambdaProxyAPI-Boto3"
LAMBDA_FUNCTION_NAME = "ApiGatewayHelloWorld-Boto3"
IAM_ROLE_NAME = "ApiGatewayLambdaRole-Boto3"
LAMBDA_CODE_FILE = "hello_world_lambda.py"
ZIP_FILE = "function.zip"
STAGE_NAME = "dev"

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

def create_iam_role():
    """Creates the IAM role for the Lambda function."""
    print("Creating IAM Role for Lambda...")
    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 API Gateway backend Lambda"
        )
        role_arn = role_response['Role']['Arn']
        print(f"IAM Role created with ARN: {role_arn}")

        iam_client.attach_role_policy(
            RoleName=IAM_ROLE_NAME,
            PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        )
        print("Attached AWSLambdaBasicExecutionRole policy.")
        
        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.")
            return iam_client.get_role(RoleName=IAM_ROLE_NAME)['Role']['Arn']
        else:
            raise e

def create_lambda_function(role_arn):
    """Packages and creates the Lambda function."""
    print("Packaging and creating Lambda function...")
    try:
        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}
        )
        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 create_api_gateway(lambda_arn):
    """Creates and configures the API Gateway."""
    print("Creating and configuring API Gateway...")
    try:
        # Create REST API
        api_response = apigateway_client.create_rest_api(name=API_NAME)
        api_id = api_response['id']
        print(f"API Gateway created with ID: {api_id}")

        # Get root resource ID
        root_resource_id = apigateway_client.get_resources(restApiId=api_id)['items'][0]['id']

        # Create a resource (e.g., /hello)
        resource_response = apigateway_client.create_resource(
            restApiId=api_id,
            parentId=root_resource_id,
            pathPart='hello'
        )
        resource_id = resource_response['id']

        # Create a GET method
        apigateway_client.put_method(
            restApiId=api_id,
            resourceId=resource_id,
            httpMethod='GET',
            authorizationType='NONE'
        )

        # Set up Lambda proxy integration
        integration_uri = f"arn:aws:apigateway:{REGION}:lambda:path/2015-03-31/functions/{lambda_arn}/invocations"
        apigateway_client.put_integration(
            restApiId=api_id,
            resourceId=resource_id,
            httpMethod='GET',
            type='AWS_PROXY',
            integrationHttpMethod='POST',
            uri=integration_uri
        )

        # Grant API Gateway permission to invoke Lambda
        account_id = sts_client.get_caller_identity()['Account']
        source_arn = f"arn:aws:execute-api:{REGION}:{account_id}:{api_id}/*/*/*"
        lambda_client.add_permission(
            FunctionName=LAMBDA_FUNCTION_NAME,
            StatementId=f"APIGatewayInvokePermission-{uuid.uuid4().hex}",
            Action='lambda:InvokeFunction',
            Principal='apigateway.amazonaws.com',
            SourceArn=source_arn
        )

        # Deploy the API
        apigateway_client.create_deployment(restApiId=api_id, stageName=STAGE_NAME)
        print(f"API deployed to stage '{STAGE_NAME}'.")

        api_url = f"https://{api_id}.execute-api.{REGION}.amazonaws.com/{STAGE_NAME}/hello"
        return api_url

    except ClientError as e:
        print(f"Error creating API Gateway: {e}")
        raise

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

        print("\n--- API Gateway setup complete! ---")
        print(f"API Endpoint URL: {api_url}")

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

if __name__ == "__main__":
    main()