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

# A script to create a simple Step Functions state machine using Boto3.

# --- Configuration ---
REGION = "us-east-1"
LAMBDA_FUNCTION_NAME = "StepFunctionsHelloWorldBoto3"
LAMBDA_CODE_FILE = "hello_world_lambda.py"
ZIP_FILE = "function.zip"
SFN_ROLE_NAME = "StepFunctionsExecutionRoleBoto3"
STATE_MACHINE_NAME = "MyBoto3SimpleStateMachine"

iam_client = boto3.client('iam', region_name=REGION)
lambda_client = boto3.client('lambda', region_name=REGION)
sfn_client = boto3.client('stepfunctions', region_name=REGION)

def create_lambda_function():
    """Creates a simple "Hello World" Lambda function."""
    print(f"--- Creating Lambda Function: {LAMBDA_FUNCTION_NAME} ---")

    # Create dummy Lambda code
    with open(LAMBDA_CODE_FILE, "w") as f:
        f.write("""
import json

def lambda_handler(event, context):
    print(f"Received event: {json.dumps(event)}")
    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Step Functions Lambda!')
    }
""")

    # Zip the 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()

    # Create IAM Role for Lambda
    lambda_role_name = f"{LAMBDA_FUNCTION_NAME}Role"
    lambda_trust_policy = {
      "Version": "2012-10-17",
      "Statement": [{"Effect": "Allow", "Principal": {"Service": "lambda.amazonaws.com"}, "Action": "sts:AssumeRole"}]
    }
    try:
        lambda_role_response = iam_client.create_role(
            RoleName=lambda_role_name,
            AssumeRolePolicyDocument=json.dumps(lambda_trust_policy)
        )
        lambda_role_arn = lambda_role_response['Role']['Arn']
        iam_client.attach_role_policy(
            RoleName=lambda_role_name,
            PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        )
        print("Waiting for Lambda IAM role to propagate...")
        time.sleep(10)
    except ClientError as e:
        if e.response['Error']['Code'] == 'EntityAlreadyExists':
            print(f"Lambda role '{lambda_role_name}' already exists. Fetching ARN.")
            lambda_role_arn = iam_client.get_role(RoleName=lambda_role_name)['Role']['Arn']
        else:
            raise e

    # Create Lambda function
    try:
        lambda_response = lambda_client.create_function(
            FunctionName=LAMBDA_FUNCTION_NAME,
            Runtime='python3.9',
            Role=lambda_role_arn,
            Handler=f"{os.path.splitext(LAMBDA_CODE_FILE)[0]}.lambda_handler",
            Code={'ZipFile': zipped_code}
        )
        lambda_arn = lambda_response['FunctionArn']
        print(f"Lambda function created with ARN: {lambda_arn}")
        return lambda_arn, lambda_role_name
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceConflictException':
            print(f"Lambda function '{LAMBDA_FUNCTION_NAME}' already exists. Fetching ARN.")
            lambda_arn = lambda_client.get_function(FunctionName=LAMBDA_FUNCTION_NAME)['Configuration']['FunctionArn']
            return lambda_arn, lambda_role_name
        else:
            print(f"Error creating Lambda function: {e}")
            raise
    finally:
        os.remove(ZIP_FILE)
        os.remove(LAMBDA_CODE_FILE)

def create_sfn_iam_role(lambda_arn):
    """Creates an IAM role for the Step Functions state machine."""
    print(f"\n--- Creating IAM Role for Step Functions: {SFN_ROLE_NAME} ---")
    sfn_trust_policy = {
      "Version": "2012-10-17",
      "Statement": [{"Effect": "Allow", "Principal": {"Service": "states.amazonaws.com"}, "Action": "sts:AssumeRole"}]
    }
    try:
        sfn_role_response = iam_client.create_role(
            RoleName=SFN_ROLE_NAME,
            AssumeRolePolicyDocument=json.dumps(sfn_trust_policy)
        )
        sfn_role_arn = sfn_role_response['Role']['Arn']
        
        sfn_policy_document = {
          "Version": "2012-10-17",
          "Statement": [
            {"Effect": "Allow", "Action": ["lambda:InvokeFunction"], "Resource": lambda_arn},
            {"Effect": "Allow", "Action": ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"], "Resource": "arn:aws:logs:*:*:*"} # For logging
          ]
        }
        sfn_policy_response = iam_client.create_policy(
            PolicyName=f"{SFN_ROLE_NAME}Policy",
            PolicyDocument=json.dumps(sfn_policy_document)
        )
        sfn_policy_arn = sfn_policy_response['Policy']['Arn']
        iam_client.attach_role_policy(RoleName=SFN_ROLE_NAME, PolicyArn=sfn_policy_arn)
        print("Waiting for SFN IAM role to propagate...")
        time.sleep(10)
        return sfn_role_arn, sfn_policy_arn
    except ClientError as e:
        if e.response['Error']['Code'] == 'EntityAlreadyExists':
            print(f"SFN role '{SFN_ROLE_NAME}' already exists. Fetching ARN.")
            sfn_role_arn = iam_client.get_role(RoleName=SFN_ROLE_NAME)['Role']['Arn']
            sfn_policy_arn = iam_client.list_policies(Scope='Local', PolicyFilter='Local')['Policies'][0]['Arn'] # This is a simplification, better to store policy ARN
            return sfn_role_arn, sfn_policy_arn
        else:
            raise e

def create_state_machine(sfn_role_arn, lambda_arn):
    """Creates a Step Functions state machine."""
    print(f"\n--- Creating State Machine: {STATE_MACHINE_NAME} ---")
    state_machine_definition = {
      "Comment": "A simple state machine that invokes a Lambda function.",
      "StartAt": "InvokeLambda",
      "States": {
        "InvokeLambda": {
          "Type": "Task",
          "Resource": "arn:aws:states:::lambda:invoke",
          "Parameters": {
            "FunctionName": lambda_arn,
            "Payload.$": "$"
          },
          "End": True
        }
      }
    }
    try:
        response = sfn_client.create_state_machine(
            name=STATE_MACHINE_NAME,
            definition=json.dumps(state_machine_definition),
            roleArn=sfn_role_arn,
            type='STANDARD'
        )
        state_machine_arn = response['stateMachineArn']
        print(f"State Machine created with ARN: {state_machine_arn}")
        return state_machine_arn
    except ClientError as e:
        if e.response['Error']['Code'] == 'StateMachineAlreadyExists':
            print(f"State Machine '{STATE_MACHINE_NAME}' already exists. Fetching ARN.")
            response = sfn_client.list_state_machines(maxResults=100)
            for sm in response['stateMachines']:
                if sm['name'] == STATE_MACHINE_NAME:
                    return sm['stateMachineArn']
            raise Exception("Could not find existing state machine.")
        else:
            print(f"Error creating state machine: {e}")
            raise

def cleanup_resources(lambda_arn, lambda_role_name, sfn_arn, sfn_role_arn, sfn_policy_arn):
    """Cleans up all created resources."""
    print(f"\n--- Cleaning up resources ---")

    # Delete State Machine
    if sfn_arn:
        print(f"Deleting State Machine '{STATE_MACHINE_NAME}'...")
        try:
            sfn_client.delete_state_machine(stateMachineArn=sfn_arn)
            print("State Machine deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'StateMachineDoesNotExist':
                print(f"State Machine '{STATE_MACHINE_NAME}' not found, skipping deletion.")
            else:
                print(f"Error deleting state machine: {e}")

    # Detach and Delete SFN IAM Role
    if sfn_role_arn:
        print(f"Detaching policy from SFN IAM Role '{SFN_ROLE_NAME}'...")
        try:
            iam_client.detach_role_policy(RoleName=SFN_ROLE_NAME, PolicyArn=sfn_policy_arn)
            iam_client.delete_policy(PolicyArn=sfn_policy_arn)
            print(f"Deleting SFN IAM Role '{SFN_ROLE_NAME}'...")
            iam_client.delete_role(RoleName=SFN_ROLE_NAME)
            print("Step Functions IAM Role deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchEntity':
                print(f"SFN Role '{SFN_ROLE_NAME}' not found, skipping deletion.")
            else:
                print(f"Error deleting SFN role: {e}")

    # Delete Lambda Function
    if lambda_arn:
        print(f"Deleting Lambda Function '{LAMBDA_FUNCTION_NAME}'...")
        try:
            lambda_client.delete_function(FunctionName=LAMBDA_FUNCTION_NAME)
            print("Lambda Function deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'ResourceNotFoundException':
                print(f"Lambda Function '{LAMBDA_FUNCTION_NAME}' not found, skipping deletion.")
            else:
                print(f"Error deleting Lambda function: {e}")

    # Detach and Delete Lambda IAM Role
    if lambda_role_name:
        print(f"Detaching policy from Lambda IAM Role '{lambda_role_name}'...")
        try:
            iam_client.detach_role_policy(RoleName=lambda_role_name, PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole')
            print(f"Deleting Lambda IAM Role '{lambda_role_name}'...")
            iam_client.delete_role(RoleName=lambda_role_name)
            print("Lambda IAM Role deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchEntity':
                print(f"Lambda Role '{lambda_role_name}' not found, skipping deletion.")
            else:
                print(f"Error deleting Lambda role: {e}")

def main():
    lambda_arn = None
    lambda_role_name = None
    sfn_role_arn = None
    sfn_policy_arn = None
    sfn_arn = None
    try:
        lambda_arn, lambda_role_name = create_lambda_function()
        sfn_role_arn, sfn_policy_arn = create_sfn_iam_role(lambda_arn)
        sfn_arn = create_state_machine(sfn_role_arn, lambda_arn)

        print("\n--- Step Functions State Machine Setup Complete! ---")
        print(f"State Machine ARN: {sfn_arn}")
        print("You can now start an execution of this state machine.")

        input("Press Enter to delete the state machine and clean up resources...")

    except ClientError as e:
        print(f"An AWS client error occurred: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
    finally:
        cleanup_resources(lambda_arn, lambda_role_name, sfn_arn, sfn_role_arn, sfn_policy_arn)
        print("\n--- Step Functions demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()