#!/bin/bash

# Function to display usage
usage() {
    echo "Usage: $0 [list|remove|list-all-extensions] <extension-arn>"
    echo "  list               : List all Lambda functions with the specified extension"
    echo "  remove             : Remove the specified extension from all Lambda functions"
    echo "  list-all-extensions: List all extensions with their ARNs and associated Lambda functions"
    echo "Note: <extension-arn> is not required for list-all-extensions option"
    exit 1
}

# Check if correct number of arguments are provided
if [ "$1" != "list-all-extensions" ] && [ $# -ne 2 ]; then
    usage
fi

ACTION=$1
EXTENSION_ARN=$2

# Function to get all Lambda functions
get_functions() {
    aws lambda list-functions --query 'Functions[*].FunctionName' --output text
}

# Function to check if a function has the specified extension
has_extension() {
    local function_name=$1
    aws lambda get-function-configuration --function-name "$function_name" --query "Layers[?Arn=='$EXTENSION_ARN'].Arn" --output text | grep -q "$EXTENSION_ARN"
}

# Function to remove the extension from a function
remove_extension() {
    local function_name=$1
    local layers=$(aws lambda get-function-configuration --function-name "$function_name" --query "Layers[?Arn!='$EXTENSION_ARN'].Arn" --output json)
    aws lambda update-function-configuration --function-name "$function_name" --layers "$layers"
}

# Function to list all extensions
list_all_extensions() {
    for func in $(get_functions); do
        echo "Function: $func"
        aws lambda get-function-configuration --function-name "$func" --query 'Layers[*].[Arn]' --output text | while read -r layer_arn; do
            echo "  Extension ARN: $layer_arn"
        done
        echo ""
    done
}

# Main logic
case $ACTION in
    list)
        echo "Listing Lambda functions with extension: $EXTENSION_ARN"
        echo "---------------------------------------------"
        for func in $(get_functions); do
            if has_extension "$func"; then
                echo "$func"
            fi
        done
        echo "---------------------------------------------"
        ;;
    remove)
        echo "Removing extension $EXTENSION_ARN from Lambda functions"
        echo "---------------------------------------------"
        for func in $(get_functions); do
            if has_extension "$func"; then
                echo "Removing from $func"
                remove_extension "$func"
            fi
        done
        echo "---------------------------------------------"
        ;;
    list-all-extensions)
        echo "Listing all extensions for all Lambda functions:"
        echo "---------------------------------------------"
        list_all_extensions
        echo "---------------------------------------------"
        ;;
    *)
        usage
        ;;
esac