#!/bin/bash
set -euo pipefail

START=$(date +%s)
SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
CLUSTER_NAME_BASE="nth-test"
PRESERVE=false
TMP_DIR=""
DOCKER_ARGS=""
PROVISION_CLUSTER_ARGS=""
DELETE_CLUSTER_ARGS=""
DEFAULT_NODE_TERMINATION_HANDLER_DOCKER_IMG="node-termination-handler:customtest"
NODE_TERMINATION_HANDLER_DOCKER_IMG=""
DEFAULT_EC2_METADATA_DOCKER_IMG="ec2-metadata-test-proxy:customtest"
EC2_METADATA_DOCKER_IMG=""
OVERRIDE_PATH=0
K8S_VERSION="1.16"
ASSERTION_SCRIPTS=$(find $SCRIPTPATH/../e2e -type f | sort)

function timeout() { perl -e 'alarm shift; exec @ARGV' "$@"; }

function relpath() {
  perl -e 'use File::Spec; print File::Spec->abs2rel(@ARGV) . "\n"' $1 $2
}

function clean_up {
    if [[ "$PRESERVE" == false ]]; then
        $SCRIPTPATH/../k8s-local-cluster-test/delete-cluster $DELETE_CLUSTER_ARGS || :
        return 
    fi 
    echo "To resume test with the same cluster use: \"-c $TMP_DIR\""""
}

function exit_and_fail {
    POD_ID=$(kubectl get pods -n kube-system | grep -i node-termination-handler | grep Running | cut -d' ' -f1 )
    kubectl logs $POD_ID --namespace kube-system
    END=$(date +%s)
    echo "⏰ Took $(expr $END - $START)sec"
    echo "❌ NTH Integration Test FAILED $CLUSTER_NAME! ❌"
    exit 1
}

function reset_cluster {
    echo "Resetting cluster"
    charts=$(helm ls --all --short)
    if [[ ! -z "$charts" ]]; then 
        helm del $charts || :
    fi 
    system_charts=$(helm ls --all --short --namespace kube-system)
    if [[ ! -z "$system_charts" ]]; then 
        helm del $system_charts --namespace kube-system || :
    fi
    for node in $(kubectl get nodes | tail -n+2 | cut -d' ' -f1); do 
        kubectl uncordon $node
        kubectl taint node $node aws-node-termination-handler/scheduled-maintenance- || true
        kubectl taint node $node aws-node-termination-handler/spot-itn- || true
    done
    remove_labels || :   
    sleep 2
}

function remove_labels {
    echo "Removing labels from NTH cluster nodes"
    labels_to_remove=($(kubectl get nodes -o json | jq '.items[].metadata.labels' | grep 'aws-node-termination-handler' | tr -d '[:blank:]' | tr -d '\"' | cut -d':' -f1))
    if [[ "${#labels_to_remove[@]}" -ne 0 ]]; then
      for l in "${labels_to_remove[@]}"; do 
        for n in $(kubectl get nodes -o json | jq -r '.items[].metadata.name'); do
          echo "Deleting label $l on node $n"
          kubectl label node $n "$l"-
        done
      done
    fi
}

function get_nth_worker_pod {
    kubectl get pods -n kube-system \
      --selector 'app.kubernetes.io/name=aws-node-termination-handler' \
      --field-selector="spec.nodeName=$CLUSTER_NAME-worker,status.phase=Running" \
      --sort-by=.metadata.creationTimestamp \
      --output jsonpath='{.items[-1].metadata.name}'
}

USAGE=$(cat << 'EOM'
  Usage: run-test [-p] [-d] [-o] [-a <ASSERTION_SCRIPT] [-b <TEST_BASE_NAME] [-c <CLUSTER_CONTEXT_DIR] [-n <NTH_DOCKER_IMG>] [-e <IMDS_DOCKER_IMG>] [-v K8S_VERSION]
  Executes a test within a provisioned kubernetes cluster with NTH and IMDS pre-loaded.

  Example: run-test -p -n node-termination-handler:customtest

          Optional:
            -a          Assertion script (default is ALL tests in the e2e dir) 
            -b          Base name of test (will be used for cluster too)
            -c          Cluster context directory, if operating on an existing cluster
            -p          Preserve kind k8s cluster for inspection
            -n          Node Termination Handler Docker Image
            -e          EC2 Metadata Docker Image 
            -d          use GOPROXY=direct to bypass proxy.golang.org
            -o          Override path w/ your own kubectl and kind binaries
            -v          Kubernetes Version (Default: 1.16) [1.12, 1.13, 1.14, 1.15, 1.16, 1.17, and 1.18]
            
EOM
)

# Process our input arguments
while getopts "pdn:e:oc:a:b:v:" opt; do
  case ${opt} in
    p ) # PRESERVE K8s Cluster
        echo "❄️  This run will preserve the cluster as you requested"
        PRESERVE=true
      ;;
    n ) # Node Termination Handler Docker Image 
        NODE_TERMINATION_HANDLER_DOCKER_IMG="$OPTARG"
      ;;
    e ) # EC2 Metadata Docker Image
        EC2_METADATA_DOCKER_IMG="$OPTARG"
      ;;
    d ) # use GOPROXY=direct
        DOCKER_ARGS="--build-arg GOPROXY=direct"
      ;;
    o ) # Override path with your own kubectl and kind binaries
        DELETE_CLUSTER_ARGS="$DELETE_CLUSTER_ARGS -o"
        PROVISION_CLUSTER_ARGS="$PROVISION_CLUSTER_ARGS -o"
        OVERRIDE_PATH=1
      ;;
    c ) # Cluster context directory to operate on existing cluster
        TMP_DIR=$OPTARG
      ;;
    a ) # Assertion script
        ASSERTION_SCRIPTS=$(echo $OPTARG | tr "," "\n")
      ;;
    b ) # Base cluster name
        CLUSTER_NAME_BASE=$OPTARG
      ;;
    v ) # K8s VERSION
        K8S_VERSION=$OPTARG
      ;;
    \? )
        echo "$USAGE" 1>&2
        exit
      ;;
  esac
done

if [[ -z $TMP_DIR ]]; then 
    TMP_DIR=$($SCRIPTPATH/../k8s-local-cluster-test/provision-cluster -b $CLUSTER_NAME_BASE -v $K8S_VERSION $PROVISION_CLUSTER_ARGS)
fi
if [[ $OVERRIDE_PATH -eq 0 ]]; then 
  export PATH=$TMP_DIR:$PATH
else 
  export PATH=$PATH:$TMP_DIR
fi
CLUSTER_NAME=$(cat $TMP_DIR/clustername)

## Build and Load Docker Images

if [ -z "$NODE_TERMINATION_HANDLER_DOCKER_IMG" ]; then 
    echo "🥑 Building the node-termination-handler docker image"
    docker build $DOCKER_ARGS -t $DEFAULT_NODE_TERMINATION_HANDLER_DOCKER_IMG "$SCRIPTPATH/../../." 
    NODE_TERMINATION_HANDLER_DOCKER_IMG="$DEFAULT_NODE_TERMINATION_HANDLER_DOCKER_IMG"
    echo "👍 Built the node-termination-handler docker image"
else 
    echo "🥑 Skipping building the node-termination-handler docker image, since one was specified ($NODE_TERMINATION_HANDLER_DOCKER_IMG)"
fi
echo "$NODE_TERMINATION_HANDLER_DOCKER_IMG" > $TMP_DIR/nth-docker-img

if [ -z "$EC2_METADATA_DOCKER_IMG" ]; then 
    echo "🥑 Building the ec2-metadata-test-proxy docker image"
    docker build $DOCKER_ARGS -t $DEFAULT_EC2_METADATA_DOCKER_IMG "$SCRIPTPATH/../ec2-metadata-test-proxy/."
    EC2_METADATA_DOCKER_IMG="$DEFAULT_EC2_METADATA_DOCKER_IMG"
    echo "👍 Built the ec2-metadata-test-proxy docker image"
else  
    echo "🥑 Skipping building the ec2-metadata-test-proxy docker image, since one was specified ($EC2_METADATA_DOCKER_IMG)"
fi
echo "$EC2_METADATA_DOCKER_IMG" > $TMP_DIR/ec2-metadata-test-proxy-docker-img

echo "🥑 Loading both images into the cluster"
kind load docker-image --name $CLUSTER_NAME --nodes=$CLUSTER_NAME-worker,$CLUSTER_NAME-control-plane $NODE_TERMINATION_HANDLER_DOCKER_IMG
kind load docker-image --name $CLUSTER_NAME --nodes=$CLUSTER_NAME-worker,$CLUSTER_NAME-control-plane $EC2_METADATA_DOCKER_IMG
echo "👍 Loaded both images into the cluster"

export KUBECONFIG="$TMP_DIR/kubeconfig"

trap "exit_and_fail" INT TERM ERR
trap "clean_up" EXIT

echo "======================================================================================================"
echo "To poke around your test manually:"
echo "export KUBECONFIG=$TMP_DIR/kubeconfig"
echo "export PATH=$TMP_DIR:\$PATH"
echo "kubectl get pods -A"
echo "======================================================================================================"


### exported vars and funcs that tests can use
export TMP_DIR
export CLUSTER_NAME
export -f timeout
export -f relpath
export -f get_nth_worker_pod
export NODE_TERMINATION_HANDLER_DOCKER_IMG=$(cat $TMP_DIR/nth-docker-img)
export NODE_TERMINATION_HANDLER_DOCKER_REPO=$(echo $NODE_TERMINATION_HANDLER_DOCKER_IMG | cut -d':' -f1)
export NODE_TERMINATION_HANDLER_DOCKER_TAG=$(echo $NODE_TERMINATION_HANDLER_DOCKER_IMG | cut -d':' -f2)
export EC2_METADATA_DOCKER_IMG=$(cat $TMP_DIR/ec2-metadata-test-proxy-docker-img)
export EC2_METADATA_DOCKER_REPO=$(echo $EC2_METADATA_DOCKER_IMG | cut -d':' -f1)
export EC2_METADATA_DOCKER_TAG=$(echo $EC2_METADATA_DOCKER_IMG | cut -d':' -f2)
###

i=0
for assert_script in $ASSERTION_SCRIPTS; do
  reset_cluster
  export IMDS_PORT=$(expr $i + 1338)
  i=$(expr $i + 1)
  echo "======================================================================================================"
  echo "🥑 Running assertion script $(basename $assert_script)"
  echo "======================================================================================================"
  assert_start=$(date +%s)
  $assert_script
  assert_end=$(date +%s)
  echo "⏰ Took $(expr $assert_end - $assert_start)sec"
  POD_ID=$(get_nth_worker_pod || :)
  kubectl logs $POD_ID --namespace kube-system || :
  ## Resets cluster to run another test on the same cluster
  reset_cluster
  echo "✅ Assertion test $assert_script PASSED! ✅"
done

echo "======================================================================================================"
echo "✅ All tests passed! ✅"
echo "======================================================================================================"

