#!/bin/bash

# IAM Role permissions needed to run this script
#   lambda:ListFunctions
#   lambda:GetFunction
#   lambda:UpdateFunctionConfiguration

# Function to update the environment variables
update_environment() {
    local FUNCTION_NAME=$1
	local REGION=$2
	local CONTROLLER=$3
	local PORT=$4
	local PROTOCOL=$5
    local WRAPPER_SCRIPT="/opt/appsentinels_wrapper.sh"

    # Set environment variables
	echo "Setting the environment variables..."
    ENV_VAR_UPDATE_OUTPUT=$(aws lambda update-function-configuration \
        --function-name $FUNCTION_NAME \
        --environment "Variables={AS_CONTROLLER_ENDPOINT=$CONTROLLER,AS_CONTROLLER_PORT=$PORT,AS_RELAY_PROTOCOL=$PROTOCOL,AWS_LAMBDA_EXEC_WRAPPER=$WRAPPER_SCRIPT}" \
		--region $REGION)
	echo "Environment variable update output: $ENV_VAR_UPDATE_OUTPUT"

	# Wait again for the function update to complete
    echo "Waiting for environment variable update to complete..."
    aws lambda wait function-updated --function-name $FUNCTION_NAME --region $REGION
    echo "Environment update complete."
}

# Function to update a single Lambda function
update_function() {
    local FUNCTION_NAME=$1
	local REGION=$2
	local LAYER_ARN=$3

    echo "Updating function: $FUNCTION_NAME"
	LAYER_NAME=$(echo $LAYER_ARN | cut -d':' -f7)
	NEW_VERSION=$(echo $LAYER_ARN | cut -d':' -f8)
    
    # Get current layers
    CURRENT_LAYERS=$(aws lambda get-function --function-name $FUNCTION_NAME --region $REGION --query 'Configuration.Layers[*].Arn' --output json)
	echo "Current layers: $CURRENT_LAYERS"

	# If CURRENT_LAYERS is null or empty, set UPDATED_LAYERS to only contain the new layer
    if [ "$CURRENT_LAYERS" == "null" ] || [ "$CURRENT_LAYERS" == "[]" ]; then
        UPDATED_LAYERS="[\"$LAYER_ARN\"]"
    else
        # Remove the old layer version and add the new one
        UPDATED_LAYERS=$(echo $CURRENT_LAYERS | jq '[.[] | select(. | contains("'$LAYER_NAME'") | not)] + ["'$LAYER_ARN'"]')
    fi
 
	echo "Updated layers: $UPDATED_LAYERS"
    
    # Update the function with the new layer list
    UPDATE_OUTPUT=$(aws lambda update-function-configuration \
        --function-name $FUNCTION_NAME \
        --layers $(echo $UPDATED_LAYERS | jq -r 'join(" ")') \
		--region $REGION)
	echo "Update output: $UPDATE_OUTPUT"

	echo "Waiting for function update to complete..."
    aws lambda wait function-updated --function-name $FUNCTION_NAME --region $REGION

	# Verify the update
    VERIFIED_LAYERS=$(aws lambda get-function --function-name $FUNCTION_NAME --region $REGION --query 'Configuration.Layers[*].Arn' --output json)
    echo "Verified layers after update: $VERIFIED_LAYERS"

	if echo $VERIFIED_LAYERS | grep -q "$LAYER_ARN"; then
        echo "Successfully updated $FUNCTION_NAME to use layer version $NEW_VERSION"
    else
        echo "Failed to update $FUNCTION_NAME to use layer version $NEW_VERSION"
    fi

    # echo "Updated $FUNCTION_NAME"
	echo "Updated $FUNCTION_NAME to use layer version $NEW_VERSION"
}

# Check if an argument is provided
if [ $# -lt 5 ]; then
    echo "Usage: $0 <python-layer-arn> <go-layer-arn> <region> <function-name | all> <controller-url>"
    exit 1
fi

# Get the arguments
PYTHONLAYER_ARN="$1"
GOLAYER_ARN="$2"
REGION="$3"
FUNCTION_ARG=$4
CONTROLLER_URL=$5
CONTROLLER_PORT="9004"
RELAY_PROTOCOL="http"

if [ "$FUNCTION_ARG" = "all" ]; then
    # Get all Lambda functions in the region
    FUNCTIONS=$(aws lambda list-functions --region $REGION --query 'Functions[*].FunctionName' --output text)
    
    # Loop through each function and update it
    for FUNCTION_NAME in $FUNCTIONS
    do
        # Update the python layer
        update_function $FUNCTION_NAME $REGION $PYTHONLAYER_ARN

        # Update the go layer
        update_function $FUNCTION_NAME $REGION $GOLAYER_ARN

        # Update the environment variables
        update_environment $FUNCTION_NAME $REGION $CONTROLLER_URL $CONTROLLER_PORT $RELAY_PROTOCOL
    done
    
    echo "All functions updated"
else
    # Check if the function exists
    if aws lambda get-function --function-name $FUNCTION_ARG --region $REGION &> /dev/null; then
        update_function $FUNCTION_ARG $REGION $PYTHONLAYER_ARN
        update_function $FUNCTION_ARG $REGION $GOLAYER_ARN
        update_environment $FUNCTION_ARG $REGION $CONTROLLER_URL $CONTROLLER_PORT $RELAY_PROTOCOL

    else
        echo "Error: Function $FUNCTION_ARG does not exist in region $REGION"
        exit 1
    fi
fi