#!/usr/bin/env bash
# Manage deployments.
#
# COPYRIGHT:   (c) Awasu Pty. Ltd. 2024 (all rights reserved).
#              Unauthorized use of this code is prohibited.
#
# LICENSE:     This software is provided 'as-is', without any express
#              or implied warranty.
#
#              In no event will the author be held liable for any damages
#              arising from the use of this software.
#
#              Permission is granted to anyone to use this software
#              for any purpose, and to alter it and redistribute it freely,
#              subject to the following restrictions:
#
#              - The origin of this software must not be misrepresented;
#                you must not claim that you wrote the original software.
#                If you use this software, an acknowledgement is requested
#                but not required.
#
#              - Altered source versions must be plainly marked as such,
#                and must not be misrepresented as being the original software.
#                Altered source is encouraged to be submitted back to
#                the original author so it can be shared with the community.
#                Please share your changes.
#
#              - This notice may not be removed or altered from any
#                source distribution.
#
# NOTE: While we can easily scale down all deployments like this:
#   kubectl scale deploy --replicas=0 --all
# we want to also wait for the pods to finish.
#
# NOTE: K3s' k3s-killall.sh script stops *everything*, including K3s :-/
#   https://docs.k3s.io/upgrades/killall
#
# NOTE: There are race conditions if you start/stop deployments while they're
# already scaling up or down, so don't do that :-/

set -euo pipefail

DEPLOYMENTS_FNAME=~/.deployments

# ---------------------------------------------------------------------

function get_deployments() {
    # get the active deployments
    kubectl get deployments -o=jsonpath='{.items[?( @.status.availableReplicas > 0 )].metadata.name}'
}

function get_pods() {
    # get the active pods
    kubectl get pods -o=jsonpath='{.items..metadata.name}'
}

# ---------------------------------------------------------------------

function stop_deployments {

    # get the currently-active deployments
    active=$( get_deployments )

    # get the deployments to be stopped
    if [ -f "$DEPLOYMENTS_FNAME" ]; then
        # just use the pre-defined list of deployment names
        dnames=$( tr '\n' ' ' <"$DEPLOYMENTS_FNAME" )
    else
        dnames=$active
        # save the deployment names
        echo "Creating the deployments file: $DEPLOYMENTS_FNAME"
        echo "$dnames" | tr ' ' '\n' | sort >"$DEPLOYMENTS_FNAME"
        echo
    fi

    # scale down all the deployments
    echo "Scaling down deployments..."
    nscaled=0
    for dname in $dnames ; do
        if echo "$active" | grep -q "\b$dname\b" ; then
            kubectl scale deployment --replicas=0 "$dname" | sed -e 's/^deployment.apps\//- /'
            nscaled=$(( nscaled + 1 ))
        fi
    done
    [ $nscaled -eq 0 ] && return
    echo

    MAX_WAIT=60
    nwait=0
    echo "Waiting for deployments ($nscaled) to scale down..."
    while true ; do

        # check if there are any pods associated with our deployments
        # NOTE: This is really slow :-/ We could assume that we want all pods in the default namespace,
        # or rely on the fact that pod names are derived from their deployment, but let's do things properly...
        pods=$(
            for pod in $( get_pods) ; do
                # get the owner ReplicaSet for the next pod
                # NOTE: ownerReferences is an array, but we assume there's only 1 entry.
                rset=$( kubectl get pod "$pod" -o=jsonpath='{.metadata.ownerReferences[0].name}' 2>/dev/null )
                [ -z "$rset" ] && continue
                # get the owner Deployment for the ReplicaSet
                deploy=$( kubectl get replicaset "$rset" -o=jsonpath='{.metadata.ownerReferences[0].name}' 2>/dev/null )
                [ -z "$deploy" ] && continue
                # check if the owning Deployment is one of ours
                if echo "$dnames" | grep -q "\b$deploy\b" ; then
                    echo "$deploy"
                fi
            done
        )
        pods=$( echo "$pods" | sort | uniq )
        [ -z "$pods" ] && break

        # yup, pods are still active - check if we've been waiting for too long
        ((++nwait))
        if [ $nwait -ge $MAX_WAIT ]; then
            echo "- Failed - deployments are still active:"
            kubectl get pods | sed -e 's/^/    /'
            return 2
        fi

        # nope - just log that we're still waiting
        if [[ $(( nwait % 5 )) == 0 ]]; then
            nwaiting=$( echo "$pods" | wc -l )
            #echo "- still waiting ($nwaiting): $( echo "$pods" | tr '\n' ' ' )"
            echo "- still waiting for $nwaiting..."
        fi

        # NOTE: We would normally sleep here for a bit, but we don't bother
        # since getting the pod info is so slow ;-/
        #sleep 1

    done
}

# ---------------------------------------------------------------------

function start_deployments {

    # get the currently-active deployments
    active=$( get_deployments )

    # scale up each registered deployment
    if [ ! -f "$DEPLOYMENTS_FNAME" ]; then
        echo "Can't find the deployments file: $DEPLOYMENTS_FNAME"
        return 2
    fi
    echo "Scaling up deployments..."
    while read -r dname ; do
        if echo "$active" | grep -q "\b$dname\b" ; then
            true #echo "$dname is already active."
        else
            kubectl scale deployment --replicas=1 "$dname" | sed -e 's/^deployment.apps\//- /'
        fi
    done <"$DEPLOYMENTS_FNAME"
    echo

    # wait for the deployments to become ready
    expected=$( sort <"$DEPLOYMENTS_FNAME" )
    nexpected=$( wc -l <"$DEPLOYMENTS_FNAME" )
    echo "Waiting for deployments to become ready ($nexpected)..."
    MAX_WAIT=60
    nwait=0
    while true; do

        # check if all the deployments are ready
        dnames=$( get_deployments )
        curr=$( echo "$dnames" | tr ' ' '\n' | sort )
        [ "$curr" = "$expected" ] && break

        # nope - check if we've been waiting for too long
        ((++nwait))
        if [ $nwait -ge $MAX_WAIT ]; then
            echo "- Failed - some deployments are still not ready:"
            kubectl get pods | sed -e 's/^/    /'
            return 2
        fi

        # nope - just log that we're still waiting
        if [[ $(( nwait % 5 )) == 0 ]]; then
            ncurr=$( echo "$curr" | wc -l )
            nwaiting=$(( nexpected - ncurr ))
            echo "- still waiting for $nwaiting..."
        fi
        sleep 1
    done
}

# ---------------------------------------------------------------------

function print_help {
    echo "`basename "$0"` [start|stop|restart]"
    echo "  Manage K3s deployments."
    echo
    echo "Note that a list of deployment names is stored in $DEPLOYMENTS_FNAME."
}

# ---------------------------------------------------------------------

# parse the command-line arguments
if [ $# -eq 0 ]; then
    print_help
    exit 0
fi
cmd="$1"
if [ "$cmd" = "stop" ]; then
    stop_deployments
elif [ "$cmd" = "start" ]; then
    start_deployments
elif [ "$cmd" = "restart" ]; then
    stop_deployments && echo ; start_deployments
else
    echo "Unknown command: $cmd"
    exit 1
fi