#!/bin/bash
set -euo pipefail

# Available env vars:
#   $TMP_DIR
#   $CLUSTER_NAME
#   $KUBECONFIG
#   $NODE_TERMINATION_HANDLER_DOCKER_REPO
#   $NODE_TERMINATION_HANDLER_DOCKER_TAG
#   $WEBHOOK_DOCKER_REPO
#   $WEBHOOK_DOCKER_TAG
#   $AEMM_URL
#   $AEMM_VERSION

SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
SQUID_DOCKERHUB_IMG="sameersbn/squid:3.5.27-2@sha256:e98299069f0c6e3d9b9188903518e2f44ac36b1fa5007e879af518e1c0a234af"
SQUID_DOCKER_IMG="squid:customtest"
SQUID_PORT="3128"
SQUID_URL="tcp://squid.default.svc.cluster.local:$SQUID_PORT"

function fail_and_exit {
    echo "❌ Webhook HTTP Proxy Test failed $CLUSTER_NAME ❌"
    exit "${1:-1}"
}

function get_squid_worker_pod() {
  kubectl get pods -o json | \
    jq '.items[] | select( .metadata.name | contains("squid") ) | .metadata.name as $name | select( .spec.nodeName | contains("worker") ) | .spec.nodeName as $nodename | $name' -r 2> /dev/null || echo ""
}

function start_squid() {
  kubectl delete configmap squid-config || :
  kubectl create configmap squid-config --from-file="${SCRIPTPATH}/../assets/squid.conf"

  old_squid_pods=""
  for i in $(seq 1 10); do
    echo "Checking if squid http-proxy has been terminated from previous run..."
    old_squid_pods="$(get_squid_worker_pod)"
    if [[ -n "${old_squid_pods}" ]]; then
      echo "Still waiting on squid to terminate from last run... Check $i/10"
      sleep 10
    else
      break
    fi
  done

  helm upgrade --install --namespace default "${CLUSTER_NAME}-squid" "${SCRIPTPATH}/../../config/helm/squid/" \
    --set squid.configMap="squid-config" \
    --set squid.image.repository="squid" \
    --set squid.image.tag="customtest" \
    --force \
    --wait \

  ## Squid can take a while to start, try to get a squid worker pod, if not, hope for the best when assertions are checked
  squid_worker_pods=""
  for i in $(seq 1 10); do
    echo "Checking if squid http-proxy has started..."
    squid_worker_pods="$(get_squid_worker_pod)"
    if [[ -z "${squid_worker_pods}" ]]; then
      echo "Still waiting on squid... Check $i/10"
      sleep 10
    else
      break
    fi
  done
}

echo "Starting Webhook HTTP Proxy Test for Node Termination Handler"

docker pull "$SQUID_DOCKERHUB_IMG"
docker tag "$SQUID_DOCKERHUB_IMG" "$SQUID_DOCKER_IMG"
kind load docker-image --name "$CLUSTER_NAME" --nodes="$CLUSTER_NAME-worker,$CLUSTER_NAME-control-plane" "$SQUID_DOCKER_IMG"

## Starts squid and waits for pods to be ready
start_squid
squid_worker_pods=$(get_squid_worker_pod)

common_helm_args=()
[[ "${TEST_WINDOWS-}" == "true" ]] && common_helm_args+=(--set targetNodeOs="windows")
[[ -n "${NTH_WORKER_LABEL-}" ]] && common_helm_args+=(--set nodeSelector."$NTH_WORKER_LABEL")

aemm_helm_args=(
  upgrade
  --install
  --namespace default
  "$CLUSTER_NAME-aemm"
  "$AEMM_DL_URL"
  --set servicePort="$IMDS_PORT"
  --wait
)
[[ ${#common_helm_args[@]} -gt 0 ]] &&
    aemm_helm_args+=("${common_helm_args[@]}")

set -x
retry 5 helm "${aemm_helm_args[@]}"
set +x

emtp_helm_args=(
  upgrade
  --install
  --namespace default
  "$CLUSTER_NAME-emtp"
  "$SCRIPTPATH/../../config/helm/webhook-test-proxy/"
  --set webhookTestProxy.image.repository="$WEBHOOK_DOCKER_REPO"
  --set webhookTestProxy.image.tag="$WEBHOOK_DOCKER_TAG"
  --wait
)
[[ -n "${WEBHOOK_DOCKER_PULL_POLICY-}" ]] &&
    emtp_helm_args+=(--set webhookTestProxy.image.pullPolicy="$WEBHOOK_DOCKER_PULL_POLICY")
[[ ${#common_helm_args[@]} -gt 0 ]] &&
    emtp_helm_args+=("${common_helm_args[@]}")

set -x
helm "${emtp_helm_args[@]}"
set +x

anth_helm_args=(
  upgrade
  --install
  --namespace kube-system
  "$CLUSTER_NAME-anth"
  "$SCRIPTPATH/../../config/helm/aws-node-termination-handler/"
  --set instanceMetadataURL="http://$AEMM_URL:$IMDS_PORT"
  --set image.repository="$NODE_TERMINATION_HANDLER_DOCKER_REPO"
  --set image.tag="$NODE_TERMINATION_HANDLER_DOCKER_TAG"
  --set enableSpotInterruptionDraining="true"
  --set enableScheduledEventDraining="true"
  --set webhookURL="$WEBHOOK_URL"
  --set webhookTemplate="\{\"Content\":\"[NTH][Instance Interruption] InstanceId: \{\{ \.InstanceID \}\} - Node: \{\{ \.NodeName \}\} - InstanceType: \{\{ \.InstanceType \}\} - Kind: \{\{ \.Kind \}\} - Start Time: \{\{ \.StartTime \}\}\"\}"
  --set webhookProxy="$SQUID_URL"
  --force
  --wait
)
[[ -n "${NODE_TERMINATION_HANDLER_DOCKER_PULL_POLICY-}" ]] &&
    anth_helm_args+=(--set image.pullPolicy="$NODE_TERMINATION_HANDLER_DOCKER_PULL_POLICY")
[[ ${#common_helm_args[@]} -gt 0 ]] &&
    anth_helm_args+=("${common_helm_args[@]}")

set -x
helm "${anth_helm_args[@]}"
set +x

TAINT_CHECK_CYCLES=15
TAINT_CHECK_SLEEP=15

deployed=0
for i in $(seq 1 $TAINT_CHECK_CYCLES); do
    if [[ $(kubectl get deployments regular-pod-test -o jsonpath='{.status.unavailableReplicas}') -eq 0 ]]; then
        echo "✅ Verified regular-pod-test pod was scheduled and started!"
        deployed=1
        break
    fi
      echo "Setup Loop $i/$TAINT_CHECK_CYCLES, sleeping for $TAINT_CHECK_SLEEP seconds"
    sleep $TAINT_CHECK_SLEEP
done

if [[ $deployed -eq 0 ]]; then
    echo "❌ regular-pod-test pod deployment failed"
    fail_end_exit 2
fi

cordoned=0
evicted=0
sent=0
for i in $(seq 1 $TAINT_CHECK_CYCLES); do
  if [[ $cordoned -eq 0 ]] && kubectl get nodes "${CLUSTER_NAME}-worker" | grep SchedulingDisabled; then
      echo "✅ Verified the worker node was cordoned!"
      cordoned=1
  fi

  if [[ $cordoned -eq 1 && $evicted -eq 0 && $(kubectl get deployments regular-pod-test -o=jsonpath='{.status.unavailableReplicas}') -eq 1 ]]; then
      echo "✅ Verified the regular-pod-test pod was evicted!"
      evicted=1
      pod_id=$(get_nth_worker_pod)
  fi

  if [[ $evicted -eq 1 && $sent -eq 0 ]] && kubectl logs "${pod_id}" -n kube-system | grep 'Webhook Success' >/dev/null; then
    echo "✅ Verified that webhook successfully sent"
    sent=1
  fi

  webhook_hostname=$(echo "${WEBHOOK_URL}" | sed -e 's@^[^/]*//@@' -e 's@/.*$@@')
  if [[ $sent -eq 1 ]] && kubectl exec -i "$(echo "${squid_worker_pods}" | cut -d' ' -f1)" -- cat /var/log/squid/access.log | grep "${webhook_hostname}" >/dev/null; then
     echo "✅ Verified the webhook POST used the http proxy"
        exit 0
  fi
  echo "Assertion Loop $i/$TAINT_CHECK_CYCLES, sleeping for $TAINT_CHECK_SLEEP seconds"
  sleep $TAINT_CHECK_SLEEP
done

if [[ $cordoned -eq 0 ]]; then
    echo "❌ Worker node was not cordoned"
elif [[ $evicted -eq 0 ]]; then
    echo "❌ regular-pod-test pod was not evicted"
elif [[ $sent -eq 0 ]]; then
    echo "❌ Webhook not sent"
else
    echo "❌ Webhook POST did not use http proxy"
fi

echo "=====================================   SQUID LOGS   ===================================================="
kubectl exec -i "$(echo "${squid_worker_pods}" | cut -d' ' -f1)" -- cat /var/log/squid/access.log
echo "===================================== END SQUID LOGS ===================================================="

fail_and_exit 1