#!/usr/bin/env bash
set -euo pipefail

STACK_NAME="duo-agentcore-gateway-interceptor"
BASE_URL="https://dl.duosecurity.com/agentcore-gateway-lambda-interceptor"

usage() {
  cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Download Duo AgentCore Gateway Interceptor artifacts, verify checksums,
upload to S3, and create/update the CloudFormation stack.

Designed to run in AWS CloudShell or any Linux/macOS environment with
aws CLI and curl installed.

Options:
  --version              Version to deploy, e.g. v1.0.0 (default: latest)
  --bucket               S3 bucket for Lambda code (required)
  --region               AWS region (default: current region from environment)
  --secret-arn           ARN of Duo credentials secret in Secrets Manager (required)
  --kms-key-arn          ARN of KMS key for the secret (optional, if using CMK)
  --stack-name           CloudFormation stack name (default: ${STACK_NAME})
  --log-retention-days   CloudWatch log retention in days (default: 14)
  --no-logging           Disable CloudWatch log group creation
  -h|--help              Show this help message
EOF
  exit 1
}

BUCKET=""
REGION="${AWS_DEFAULT_REGION:-${AWS_REGION:-}}"
VERSION="latest"
SECRET_ARN=""
KMS_KEY_ARN=""
LOG_RETENTION=14
CREATE_LOG_GROUP="true"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --version)            VERSION="$2"; shift 2 ;;
    --bucket)             BUCKET="$2"; shift 2 ;;
    --region)             REGION="$2"; shift 2 ;;
    --secret-arn)         SECRET_ARN="$2"; shift 2 ;;
    --kms-key-arn)        KMS_KEY_ARN="$2"; shift 2 ;;
    --stack-name)         STACK_NAME="$2"; shift 2 ;;
    --log-retention-days) LOG_RETENTION="$2"; shift 2 ;;
    --no-logging)         CREATE_LOG_GROUP="false"; shift ;;
    -h|--help)            usage ;;
    *) echo "Unknown option: $1" >&2; usage ;;
  esac
done

if [[ -z "$BUCKET" ]]; then
  echo "Error: --bucket is required." >&2
  usage
fi
if [[ -z "$SECRET_ARN" ]]; then
  echo "Error: --secret-arn is required." >&2
  usage
fi
if [[ -z "$REGION" ]]; then
  echo "Error: --region is required (or set AWS_DEFAULT_REGION)." >&2
  usage
fi

if [[ "$VERSION" == "latest" ]]; then
  URL_PATH="${BASE_URL}/latest"
  LAMBDA_ZIP="agentcore-gateway-interceptor.zip"
  LAYER_ZIP="agentcore-gateway-interceptor-duo-client-layer.zip"
  CFN_TEMPLATE="agentcore-gateway-interceptor-cfn-template.yaml"
else
  URL_PATH="${BASE_URL}/${VERSION}"
  LAMBDA_ZIP="agentcore-gateway-interceptor-${VERSION}.zip"
  LAYER_ZIP="agentcore-gateway-interceptor-duo-client-layer-${VERSION}.zip"
  CFN_TEMPLATE="agentcore-gateway-interceptor-cfn-template-${VERSION}.yaml"
fi

WORK_DIR=$(mktemp -d)
trap 'rm -rf "$WORK_DIR"' EXIT

echo "==> Downloading artifacts (version: ${VERSION})..."
for file in "$LAMBDA_ZIP" "$LAYER_ZIP" "$CFN_TEMPLATE"; do
  echo "    ${file}"
  curl -fSL --retry 3 -o "${WORK_DIR}/${file}" "${URL_PATH}/${file}"
  curl -fSL --retry 3 -o "${WORK_DIR}/${file}.sha256" "${URL_PATH}/${file}.sha256"
done

echo "==> Verifying SHA-256 checksums..."
for file in "$LAMBDA_ZIP" "$LAYER_ZIP" "$CFN_TEMPLATE"; do
  expected=$(awk '{print $1}' "${WORK_DIR}/${file}.sha256")
  if command -v sha256sum &>/dev/null; then
    actual=$(sha256sum "${WORK_DIR}/${file}" | awk '{print $1}')
  else
    actual=$(shasum -a 256 "${WORK_DIR}/${file}" | awk '{print $1}')
  fi
  if [[ "$expected" != "$actual" ]]; then
    echo "ERROR: Checksum mismatch for ${file}" >&2
    echo "  Expected: ${expected}" >&2
    echo "  Actual:   ${actual}" >&2
    exit 1
  fi
  echo "    ${file}: OK"
done

echo "==> Uploading Lambda artifacts to s3://${BUCKET}/${VERSION}/..."
aws s3 cp "${WORK_DIR}/${LAMBDA_ZIP}" "s3://${BUCKET}/${VERSION}/${LAMBDA_ZIP}" --region "$REGION"
aws s3 cp "${WORK_DIR}/${LAYER_ZIP}" "s3://${BUCKET}/${VERSION}/${LAYER_ZIP}" --region "$REGION"

echo "==> Deploying CloudFormation stack: ${STACK_NAME}..."

CFN_PARAMS=(
  "ParameterKey=DuoCredentialsSecretArn,ParameterValue=${SECRET_ARN}"
  "ParameterKey=S3Bucket,ParameterValue=${BUCKET}"
  "ParameterKey=LambdaKey,ParameterValue=${VERSION}/${LAMBDA_ZIP}"
  "ParameterKey=LayerKey,ParameterValue=${VERSION}/${LAYER_ZIP}"
  "ParameterKey=LogRetentionInDays,ParameterValue=${LOG_RETENTION}"
  "ParameterKey=CreateCloudWatchLogGroup,ParameterValue=${CREATE_LOG_GROUP}"
)

if [[ -n "$KMS_KEY_ARN" ]]; then
  CFN_PARAMS+=("ParameterKey=DuoCredentialsKmsKeyArn,ParameterValue=${KMS_KEY_ARN}")
fi

STACK_EXISTS=$(aws cloudformation describe-stacks \
  --stack-name "$STACK_NAME" \
  --region "$REGION" \
  --query 'Stacks[0].StackStatus' \
  --output text 2>/dev/null || echo "DOES_NOT_EXIST")

if [[ "$STACK_EXISTS" == "DOES_NOT_EXIST" ]]; then
  echo "    Creating new stack..."
  aws cloudformation create-stack \
    --stack-name "$STACK_NAME" \
    --template-body "file://${WORK_DIR}/${CFN_TEMPLATE}" \
    --capabilities CAPABILITY_NAMED_IAM \
    --parameters "${CFN_PARAMS[@]}" \
    --region "$REGION"

  echo "    Waiting for stack creation to complete..."
  aws cloudformation wait stack-create-complete \
    --stack-name "$STACK_NAME" \
    --region "$REGION"
else
  echo "    Updating existing stack (current status: ${STACK_EXISTS})..."
  aws cloudformation update-stack \
    --stack-name "$STACK_NAME" \
    --template-body "file://${WORK_DIR}/${CFN_TEMPLATE}" \
    --capabilities CAPABILITY_NAMED_IAM \
    --parameters "${CFN_PARAMS[@]}" \
    --region "$REGION" || {
      err=$?
      if aws cloudformation describe-stacks \
           --stack-name "$STACK_NAME" \
           --region "$REGION" \
           --query 'Stacks[0].StackStatus' \
           --output text 2>/dev/null | grep -q "UPDATE_COMPLETE"; then
        echo "    No updates needed — stack is already up to date."
      else
        exit $err
      fi
    }

  echo "    Waiting for stack update to complete..."
  aws cloudformation wait stack-update-complete \
    --stack-name "$STACK_NAME" \
    --region "$REGION"
fi

echo "==> Done! Stack outputs:"
aws cloudformation describe-stacks \
  --stack-name "$STACK_NAME" \
  --region "$REGION" \
  --query 'Stacks[0].Outputs' \
  --output table