#!/bin/bash

# Ensure that your aws CLI is configured with the right set of credentials and the right region where the Comprehend endpoint is located

# Set these to configure aws CLI using environment variables if required.
# Refer docs for ways to configure the aws CLI - https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html
# export AWS_ACCESS_KEY_ID=
# export AWS_SECRET_ACCESS_KEY=
# export AWS_DEFAULT_REGION=

if [[ "$#" -lt 3 ]]; then
    echo "This script helps setup scheduled scaling for an endpoint. It takes the weekday IU (IU for Monday 4am PST to Friday 8pm PST)
     and weekend IU as arguments and sets up automatic time-based scale up scale down."
    echo "USAGE: $0 <endpointArn> <weekdayIU> <weekendIU>"
    echo "Example: ./endpoint_scaling.sh 5 1"
    exit 1
fi

ENDPOINT_ARN=$1
echo "Endpoint Arn is $ENDPOINT_ARN"
WEEKDAY_IU=$2
echo "Desired Weekday IU's is $WEEKDAY_IU"
WEEKEND_IU=$3
echo "Desired Weekend IU's is $WEEKEND_IU"

re='^[0-9]+$'
# Validate that the Weekday and Weekend IUs are valid ints in the range (0, 10]
if ! [[ $WEEKDAY_IU =~ $re ]] || [[ $WEEKDAY_IU -le 0 ]] || [[ $WEEKDAY_IU -gt 10 ]] ; then
    echo "Invalid argument for WEEKDAY_IU. Valid values are 0 < IU <= 10."
    exit 1
fi

if ! [[ $WEEKEND_IU =~ $re ]] || [[ $WEEKEND_IU -le 0 ]] || [[ $WEEKEND_IU -gt 10 ]] ; then
    echo "Invalid argument for WEEKEND_IU. Valid values are 0 < IU <= 10."
    exit 1
fi

# Call describe on the endpoint to ensure it is a valid IN_SERVICE endpoint
STATUS=$(aws comprehend describe-endpoint --endpoint-arn $ENDPOINT_ARN | jq -rc '.EndpointProperties' | jq -r '.Status')
echo "Status of endpoint $ENDPOINT_ARN is $STATUS"
if ! [[ ${STATUS} == 'IN_SERVICE' ]]; then
    echo "$ENDPOINT_ARN is not a valid endpoint with IN_SERVICE status."
    exit 1
fi

# Get region from the endpoint-arn
REGION="$(echo $ENDPOINT_ARN | cut -d':' -f4)"
echo "Region is $REGION"
ENDPOINT_TYPE="$(echo $ENDPOINT_ARN | cut -d':' -f6 | cut -d'/' -f1)"
echo "Endpoint type is $ENDPOINT_TYPE"

aws application-autoscaling register-scalable-target \
    --service-namespace comprehend \
    --region $REGION \
    --resource-id $ENDPOINT_ARN \
    --scalable-dimension comprehend:$ENDPOINT_TYPE:DesiredInferenceUnits \
    --min-capacity $WEEKDAY_IU \
    --max-capacity $WEEKDAY_IU  

if [[ $? != 0 ]]; then
    echo "Failed to setup scheduled scaling for $ENDPOINT_ARN"
    exit 1
fi

# Read more about cron expressions here - https://docs.aws.amazon.com/AmazonCloudWatch/latest/events/ScheduledEvents.html

# Downscale on Friday 8pm PST (Saturday 3am UTC)
aws application-autoscaling put-scheduled-action \
    --service-namespace comprehend \
    --region $REGION \
    --scheduled-action-name WeekendDownscaleScheduledAction \
    --resource-id $ENDPOINT_ARN \
    --scalable-dimension comprehend:$ENDPOINT_TYPE:DesiredInferenceUnits \
    --schedule "cron(0 3 ? * SAT *)" \
    --scalable-target-action MinCapacity=$WEEKEND_IU,MaxCapacity=$WEEKEND_IU  

if [[ $? != 0 ]]; then
    echo "Failed to setup scheduled scaling for $ENDPOINT_ARN"
    exit 1
fi

# Upscale on Monday 4am PST (Monday 11am UTC)
aws application-autoscaling put-scheduled-action \
    --service-namespace comprehend \
    --region $REGION \
    --scheduled-action-name WeekdayUpscaleScheduledAction \
    --resource-id $ENDPOINT_ARN \
    --scalable-dimension comprehend:$ENDPOINT_TYPE:DesiredInferenceUnits \
    --schedule "cron(0 11 ? * MON *)" \
    --scalable-target-action MinCapacity=$WEEKDAY_IU,MaxCapacity=$WEEKDAY_IU  

if [[ $? != 0 ]]; then
    echo "Failed to setup scheduled scaling for $ENDPOINT_ARN"
    exit 1
fi

# Can verify success using the following commands:
# aws application-autoscaling describe-scalable-targets --service-namespace comprehend
# aws application-autoscaling describe-scheduled-actions --service-namespace comprehend

echo "Set weekday IU to $WEEKDAY_IU and weekend IU to $WEEKEND_IU for endpoint $ENDPOINT_ARN"
