#!/bin/bash

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

# --- Configuration ---
REGION="us-east-1"
LAMBDA_FUNCTION_NAME="StepFunctionsHelloWorld"
LAMBDA_CODE_FILE="hello_world_lambda.py"
ZIP_FILE="function.zip"
SFN_ROLE_NAME="StepFunctionsExecutionRole"
STATE_MACHINE_NAME="MyCLISimpleStateMachine"

# --- 1. Create Lambda Function ---
echo "--- Creating Lambda Function: $LAMBDA_FUNCTION_NAME ---"

# Create dummy Lambda code
cat > $LAMBDA_CODE_FILE <<EOF
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!')
    }
EOF

# Zip the code
zip $ZIP_FILE $LAMBDA_CODE_FILE

# Create IAM Role for Lambda
LAMBDA_TRUST_POLICY_JSON=$(cat <<-EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": { "Service": "lambda.amazonaws.com" },
      "Action": "sts:AssumeRole"
    }
  ]
}
EOF
)
LAMBDA_ROLE_ARN=$(aws iam create-role \
  --role-name "${LAMBDA_FUNCTION_NAME}Role" \
  --assume-role-policy-document "$LAMBDA_TRUST_POLICY_JSON" \
  --query 'Role.Arn' --output text)
aws iam attach-role-policy \
  --role-name "${LAMBDA_FUNCTION_NAME}Role" \
  --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
sleep 10 # Give IAM time to propagate

# Create Lambda function
LAMBDA_ARN=$(aws lambda create-function \
  --function-name $LAMBDA_FUNCTION_NAME \
  --runtime python3.9 \
  --zip-file fileb://$ZIP_FILE \
  --handler hello_world_lambda.lambda_handler \
  --role $LAMBDA_ROLE_ARN \
  --query 'FunctionArn' --output text)

rm $ZIP_FILE $LAMBDA_CODE_FILE
echo "Lambda function created with ARN: $LAMBDA_ARN"

# --- 2. Create IAM Role for Step Functions ---
echo -e "\n--- Creating IAM Role for Step Functions: $SFN_ROLE_NAME ---"
SFN_TRUST_POLICY_JSON=$(cat <<-EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": { "Service": "states.amazonaws.com" },
      "Action": "sts:AssumeRole"
    }
  ]
}
EOF
)
SFN_ROLE_ARN=$(aws iam create-role \
  --role-name $SFN_ROLE_NAME \
  --assume-role-policy-document "$SFN_TRUST_POLICY_JSON" \
  --query 'Role.Arn' --output text)

# Attach policy to allow SFN to invoke Lambda
SFN_POLICY_JSON=$(cat <<-EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "lambda:InvokeFunction"
      ],
      "Resource": "$LAMBDA_ARN"
    }
  ]
}
EOF
)
SFN_POLICY_ARN=$(aws iam create-policy \
  --policy-name "${SFN_ROLE_NAME}Policy" \
  --policy-document "$SFN_POLICY_JSON" \
  --query 'Policy.Arn' --output text)
aws iam attach-role-policy \
  --role-name $SFN_ROLE_NAME \
  --policy-arn $SFN_POLICY_ARN
sleep 10 # Give IAM time to propagate

echo "Step Functions Role created with ARN: $SFN_ROLE_ARN"

# --- 3. Create State Machine ---
echo -e "\n--- Creating State Machine: $STATE_MACHINE_NAME ---"
STATE_MACHINE_DEFINITION=$(cat <<-EOF
{
  "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
    }
  }
}
EOF
)
STATE_MACHINE_ARN=$(aws stepfunctions create-state-machine \
  --name $STATE_MACHINE_NAME \
  --definition "$STATE_MACHINE_DEFINITION" \
  --role-arn $SFN_ROLE_ARN \
  --type STANDARD \
  --region $REGION \
  --query 'stateMachineArn' --output text)

echo "State Machine created with ARN: $STATE_MACHINE_ARN"

# --- 4. Output State Machine ARN ---
echo -e "\n--- Step Functions State Machine Setup Complete! ---"
echo "State Machine ARN: $STATE_MACHINE_ARN"
echo "You can now start an execution of this state machine."

read -p "Press Enter to delete the state machine and clean up resources..."

# --- Clean Up ---
echo -e "\n--- Cleaning up resources ---"

# Delete State Machine
echo "Deleting State Machine '$STATE_MACHINE_NAME'"...
aws stepfunctions delete-state-machine \
  --state-machine-arn $STATE_MACHINE_ARN \
  --region $REGION

echo "State Machine deleted."

# Detach and Delete SFN IAM Role
echo "Detaching policy from SFN IAM Role '$SFN_ROLE_NAME'"...
aws iam detach-role-policy \
  --role-name $SFN_ROLE_NAME \
  --policy-arn $SFN_POLICY_ARN \
  --region $REGION
echo "Deleting SFN IAM Policy '$SFN_POLICY_NAME'"...
aws iam delete-policy \
  --policy-arn $SFN_POLICY_ARN \
  --region $REGION
echo "Deleting SFN IAM Role '$SFN_ROLE_NAME'"...
aws iam delete-role \
  --role-name $SFN_ROLE_NAME \
  --region $REGION

echo "Step Functions IAM Role deleted."

# Delete Lambda Function
echo "Deleting Lambda Function '$LAMBDA_FUNCTION_NAME'"...
aws lambda delete-function \
  --function-name $LAMBDA_FUNCTION_NAME \
  --region $REGION

echo "Lambda Function deleted."

# Detach and Delete Lambda IAM Role
echo "Detaching policy from Lambda IAM Role '${LAMBDA_FUNCTION_NAME}Role'"...
aws iam detach-role-policy \
  --role-name "${LAMBDA_FUNCTION_NAME}Role" \
  --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole \
  --region $REGION
echo "Deleting Lambda IAM Role '${LAMBDA_FUNCTION_NAME}Role'"...
aws iam delete-role \
  --role-name "${LAMBDA_FUNCTION_NAME}Role" \
  --region $REGION

echo "Lambda IAM Role deleted."

echo -e "\n--- Step Functions demonstration and cleanup complete ---"