#!/bin/bash
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

set -euo pipefail

LAMBDA_TASK_ROOT=/var/task
TMPFILE=/tmp/data

# Graceful Shutdown
_term() {
  echo "[runtime] Received SIGTERM"
  # forward SIGTERM to child procs and exit
  kill -TERM "$PID" 2>/dev/null
  echo "[runtime] Exiting"
  exit 0
}

forward_sigterm_and_wait() {
  trap _term SIGTERM
  wait "$PID"
  trap - SIGTERM
}

echo "[runtime] Initializing..."

# Initialization - load function 
echo "[runtime] Load function from $LAMBDA_TASK_ROOT/$(echo $_HANDLER | cut -d. -f1).sh"
source $LAMBDA_TASK_ROOT/"$(echo $_HANDLER | cut -d. -f1).sh"

# Processing
while true
do
  echo "[runtime] Waiting for invocation, Get /next event from http://${AWS_LAMBDA_RUNTIME_API}/2018-06-01/runtime/invocation/next"

  HEADERS="$(mktemp)"

  # Get an event. The HTTP request will block until one is received
  curl -sS -LD "$HEADERS" -X GET "http://${AWS_LAMBDA_RUNTIME_API}/2018-06-01/runtime/invocation/next" > $TMPFILE &
  PID=$!
  forward_sigterm_and_wait

  EVENT_DATA=$(<$TMPFILE)

  echo "[runtime] Received invocation: $EVENT_DATA"

  # Extract request ID by scraping response headers received above
  REQUEST_ID=$(grep -Fi Lambda-Runtime-Aws-Request-Id "$HEADERS" | tr -d '[:space:]' | cut -d: -f2)

  echo "[runtime] Executing function handler: $_HANDLER"

  # Execute the handler function from the script
    RESPONSE=$($(echo "$_HANDLER" | cut -d. -f2) "$EVENT_DATA")
  sleep 1
  echo "[runtime] Sending invocation response: $RESPONSE to http://${AWS_LAMBDA_RUNTIME_API}/2018-06-01/runtime/invocation/$REQUEST_ID/response"

  # Send the response
  curl -sS -X POST "http://${AWS_LAMBDA_RUNTIME_API}/2018-06-01/runtime/invocation/$REQUEST_ID/response"  -d "$RESPONSE" > $TMPFILE
  PID=$!
  forward_sigterm_and_wait

  STATUS_RESP=$(<$TMPFILE)

  echo "[runtime] Runtime API response: $STATUS_RESP"
done