#!/bin/bash

set -xeuo pipefail

WEBHOOK_HOLDERPOD_TEST_NS=holderpod
WEBHOOK_FASTDISK_TEST_NS=fastdisk
SCHEDULER_TEST_NS=scheduler-test
JOINT_TEST_NS=joint-test

WEBHOOK_SCHEDULERNAME_TEST_NS=schedulername

setup() {
    kind create cluster
    kind load docker-image hub.pingcap.net/qa/test-infra-webhook:latest-amd64
    kind load docker-image hub.pingcap.net/qa/test-infra-scheduler:latest-amd64
    kubectl apply -f ./manifests/common-config

    setup_mutating_webhook_fastdisk
    setup_validating_webhook_holderpod
    e2e::__wait_for_deploy default test-infra-webhook
    setup_scheduler

    setup_mutating_webhook_schedulername
}

setup_mutating_webhook_fastdisk() {
    kubectl apply -f ./manifests/mutation/fastdisk

    kubectl proxy &
    sleep 2 ## wait proxy
    local proxy_pid=$!
    curl --header "Content-Type: application/json-patch+json" \
        --request PATCH \
        --data "[{\"op\": \"replace\", \"path\": \"/status/capacity/test-infra.pingcap.net~1fast-disks\", \"value\": \"3\"}]" \
        http://localhost:8001/api/v1/nodes/kind-control-plane/status
    kill $proxy_pid
}

setup_validating_webhook_holderpod() {
    kubectl apply -f ./manifests/validation/holderpod
}

setup_scheduler() {
    kubectl apply -f ./manifests/scheduler
    e2e::__wait_for_deploy kube-system kube-scheduler
}

setup_mutating_webhook_schedulername() {
    kubectl apply -f ./manifests/mutation/schedulername
}

cleanup() {
    kind delete cluster
}

run_test() {
    # test holder pod feature
    test_mutating_webhook_fastdisk
    test_validating_webhook_holderpod
    test_scheduler
    test_joint

    # test independent webhook
    test_mutating_webhook_schedulername
}

test_mutating_webhook_fastdisk() {
    kubectl apply -f test-data/mutation/fastdisk/pvc.yaml
    kubectl apply -f test-data/mutation/fastdisk/fast-disk.pod.yaml

    local fastdisk_request=`kubectl -n $WEBHOOK_FASTDISK_TEST_NS get pod myapp -o 'go-template={{ index (index .spec.containers 0).resources.requests "test-infra.pingcap.net/fast-disks" }}'`
    if [ $fastdisk_request -ne 1 ]; then
        echo "invalid fastdisk count"
        return 1
    fi

    kubectl delete ns $WEBHOOK_FASTDISK_TEST_NS
    e2e:__run_until_fail "kubectl get ns $WEBHOOK_FASTDISK_TEST_NS"
}

test_validating_webhook_holderpod() {
    local sts_ns=$WEBHOOK_HOLDERPOD_TEST_NS
    local sts_name="test-validating-webhook"
    local holder_pod_name="holder-$sts_name-1"

    # 1. wait for statefuleset ready, holder pod should be created
    kubectl apply -f test-data/validation/holderpod
    e2e::__wait_for_statefulset $sts_ns $sts_name 60
    kubectl -n $sts_ns get pod $holder_pod_name > /dev/null

    # 2. delete holder pod
    kubectl -n $sts_ns delete pod $holder_pod_name
    e2e:__run_until_fail "kubectl -n $sts_ns get pod $holder_pod_name" # until deleted

    # 3. scale in, holder pod should be created by validating webhook on deleting validation
    kubectl -n $sts_ns patch statefulset $sts_name -p '{"spec": {"replicas": 1}}'
    e2e:__run_until_success "kubectl -n $sts_ns get pod $holder_pod_name"

    # 4. delete ns should work
    kubectl delete ns $sts_ns
    e2e:__run_until_fail "kubectl get ns $sts_ns"
}

test_scheduler() {
    local sts_ns=$SCHEDULER_TEST_NS
    local sts_name="myapp"

    # 1. create 3 holders then a statefulset with 3 max-replicas and 2 replica
    # after statefulset ready, holer-0, holder-1 should be deleted
    local hodler_pod_1_name="holder-$sts_name-1"
    kubectl apply -f test-data/scheduler
    e2e::__wait_for_statefulset $sts_ns $sts_name
    e2e:__run_until_fail "kubectl -n $sts_ns get pod $hodler_pod_1_name"

    # 2. scale statefulset to 3 replicas, holder-2 should be deleted
    local holder_pod_2_name="holder-$sts_name-2"
    kubectl -n $sts_ns scale statefulset $sts_name --replicas 3
    e2e::__wait_for_statefulset $sts_ns $sts_name
    e2e:__run_until_fail "kubectl -n $sts_ns get pod $holder_pod_2_name"
}

test_joint() {
    local test_ns=$JOINT_TEST_NS
    local sts_name="myapp"

    # 1. create an statefulset with 2 replicas and 3 max replicas, each replica request one fast-disk
    kubectl -n $test_ns apply -f test-data/joint-test/statefuleset.yaml
    e2e::__wait_for_statefulset $test_ns $sts_name

    # 2. create a race pod that requests a fast-disk (should be blocked to the end since all fast-disks are taken by statefulset)
    kubectl -n $test_ns apply -f test-data/joint-test/resource-race.pod.yaml

    # 3. scale out statefulset to 3 replicas, it should success
    kubectl -n $test_ns scale statefulset $sts_name --replicas 3
    e2e::__wait_for_statefulset $test_ns $sts_name
}

test_mutating_webhook_schedulername() {
    kubectl apply -f test-data/mutation/schedulername

    expectedSchedulerName=`kubectl get ns $WEBHOOK_SCHEDULERNAME_TEST_NS -o json | jq -r '.metadata.annotations."test-infra.pingcap.net/schedulerName"'`
    schedulerName=`kubectl -n $WEBHOOK_SCHEDULERNAME_TEST_NS get pod myapp -o json | jq -r '.spec.schedulerName'`
    if [[ $schedulerName != $expectedSchedulerName ]]; then
        echo "invalid schedulerName"
        return 1
    fi

    kubectl delete ns $WEBHOOK_SCHEDULERNAME_TEST_NS
}

e2e:__run_until_success() {
    local cmd="$1"
    local retries="${2:-60}"
    for ((i = 0; i < retries; i++)) {
        eval $cmd > /dev/null 2>&1 && return 0
        sleep 5
    }
}

e2e:__run_until_fail() {
    local cmd="$1"
    local retries="${2:-60}"
    for ((i = 0; i < retries; i++)) {
        eval $cmd > /dev/null 2>&1 || return 0
        sleep 5
    }
}

e2e::__wait_for_deploy() {
    local ns="$1"
    local name="$2"
    local retries="${3:-60}"
    echo "info: waiting for pods of deployment $ns/$name are ready (retries: $retries, interval: 1s)"
    for ((i = 0; i < retries; i++)) {
        read a b <<<$(kubectl -n $ns get deploy/$name -ojsonpath='{.spec.replicas} {.status.readyReplicas}{"\n"}')
        if [[ "$a" -gt 0 && "$a" -eq "$b" ]]; then
            echo "info: all pods of deployment $ns/$name are ready (desired: $a, ready: $b)"
            return 0
        fi
        echo "info: pods of deployment $ns/$name (desired: $a, ready: $b)"
        sleep 5
    }
    echo "info: timed out waiting for pods of deployment $ns/$name are ready"
    return 1
}

e2e::__wait_for_statefulset() {
    local ns="$1"
    local name="$2"
    local retries="${3:-60}"
    echo "info: waiting for pods of statefulset $ns/$name are ready (retries: $retries, interval: 1s)"
    for ((i = 0; i < retries; i++)) {
        read a b <<<$(kubectl -n $ns get statefulset/$name -ojsonpath='{.spec.replicas} {.status.readyReplicas}{"\n"}')
        if [[ "$a" -gt 0 && "$a" -eq "$b" ]]; then
            echo "info: all pods of statefulset $ns/$name are ready (desired: $a, ready: $b)"
            return 0
        fi
        echo "info: pods of statefulset $ns/$name (desired: $a, ready: $b)"
        sleep 5
    }
    echo "info: timed out waiting for pods of statefulset $ns/$name are ready"
    return 1
}

CLEANUP=${CLEANUP:-}
if [ -n "$CLEANUP" ]; then
    cleanup
    exit 0
fi

setup
run_test
cleanup