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

if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root"
   exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROLLING_LEDGER="${SCRIPT_DIR}"/../../rolling/kasm-rolling-ledger
KASM_HOME=/opt/kasm/current/

show_help() {
    cat <<EOF
Usage: $0 [OPTIONS]

Options:
    Rollback: Will rollback to the direct image before the current one. Specify an image tag to rollback to a specific image.
    -r
    --rollback
    -r=<image_tag>
    --rollback=<image_tag>

    Latest: Will set Kasm to continue pulling in the latest rolling images.
    -l
    --latest

    Date: Will display the creation dates of all images.
    -d
    --date

    Update: Will update the rolling ledger to track current and previous running images.
    -u
    --update

    Cleanup: Will remove any unused image information from the rolling ledger.
    -c
    --cleanup
EOF
}

check_for_ledger() {
    if [[ ! -f "$ROLLING_LEDGER" ]]; then
        echo "ERROR: Rolling ledger not found. Cannot proceed without it!"
        exit 1
    fi
}

get_image_version() {
    echo "$1" | grep -o 'version=[^,]*' | cut -d= -f2
}

get_current_digest() {
    echo "$1" | grep -o 'current-digest=[^,]*' | cut -d= -f2
}

get_old_digest() {
    echo "$1" | grep -o 'old-digest=[^,]*' | cut -d= -f2
}

get_old_date() {
    echo "$1" | grep -o 'old-create-date=[^,]*' | cut -d= -f2
}

get_current_date() {
    echo "$1" | grep -o 'current-create-date=[^,]*' | cut -d= -f2
}

change_image_tag() {
    sed -i "s|\($1\)[^\"]*|\1$2|" ${KASM_HOME}/docker/*.yaml
}

show_image_date() {
    check_for_ledger
    OLD_IFS="$IFS"
    IFS=$'\n'
    echo "=== Currently tracking the following: ==="
    for kasm_img in $(docker images --format '{{.Repository}}' 2>/dev/null | grep kasmweb)
    do
        if current_img=$(grep "$kasm_img" "$ROLLING_LEDGER"); then
            echo "$kasm_img"
            echo -e "\tRollback Image Date: " $(get_old_date "$current_img")
            echo -e "\tRunning Image Date: " $(get_current_date "$current_img")
        fi
    done
    IFS="$OLD_IFS"
}

rollback_image_tag() { 
    check_for_ledger
    OLD_IFS="$IFS"
    IFS=$'\n'
    for kasm_img in $(docker images --format '{{.Repository}}' 2>/dev/null | grep kasmweb)
    do
        if current_img=$(grep "$kasm_img" "$ROLLING_LEDGER"); then
            old_digest=$(get_old_digest "$current_img")
            
            if [[ "$old_digest" == "<none>" ]];then continue; fi
            
            if [[ "$old_digest" != "" ]]; then
                change_image_tag "$kasm_img" "@$old_digest"
            fi
        fi
    done
    IFS="$OLD_IFS"
}

use_specific_tag() {
    OLD_IFS="$IFS"
    IFS=$'\n'
    for kasm_img in $(docker images --format '{{.Repository}}' 2>/dev/null | grep kasmweb)
    do
        change_image_tag "$kasm_img" ":$1"
    done
    IFS="$OLD_IFS"
}

use_latest() {
    check_for_ledger
    OLD_IFS="$IFS"
    IFS=$'\n'
    for kasm_img in $(docker images --format '{{.Repository}}' 2>/dev/null | grep kasmweb)
    do
        if current_img=$(grep "$kasm_img" "$ROLLING_LEDGER"); then
            version=$(get_image_version "$current_img")
            version="${version%-[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9]}"

            change_image_tag "$kasm_img" ":$version"
        fi
    done
    IFS="$OLD_IFS"
}

update_rolling() {
    touch "$ROLLING_LEDGER"
    OLD_IFS="$IFS"
    IFS=$'\n'
    for kasm_img in $(docker images --format '{{.Repository}} {{.Tag}} {{.Digest}} {{.CreatedAt}}' 2>/dev/null | grep kasmweb)
    do
        image_name=$(echo $kasm_img | awk '{print $1}')
        image_tag=$(echo $kasm_img | awk '{print $2}')
        image_digest=$(echo $kasm_img | awk '{print $3}')
        image_created=$(echo $kasm_img | awk '{print $4}')

        if grep -q "$image_name":"$image_tag"\" "${KASM_HOME}"/docker/*.yaml || grep -q "$image_digest" "${KASM_HOME}"/docker/*.yaml; then
            if current_img=$(grep "$image_name" "$ROLLING_LEDGER"); then
                old_digest=$(get_old_digest $current_img)
                current_digest=$(get_current_digest $current_img)
                current_date=$(get_current_date $current_img)

                if [[ "$image_digest" != "$current_digest" ]]; then
                    sanitzied_kasm_img=${image_name//\//\\/}
                    sed -i "/^$sanitzied_kasm_img,/ s|old-digest=[^,]*|old-digest=$current_digest|" "$ROLLING_LEDGER"
                    sed -i "/^$sanitzied_kasm_img,/ s|old-create-date=[^,]*|old-create-date=$current_date|" "$ROLLING_LEDGER"
                    sed -i "/^$sanitzied_kasm_img,/ s|current-digest=[^,]*|current-digest=$image_digest|" "$ROLLING_LEDGER"
                    sed -i "/^$sanitzied_kasm_img,/ s|current-create-date=[^,]*|current-create-date=$image_created|" "$ROLLING_LEDGER"
                fi
            else
                echo $image_name,version=$image_tag,old-digest=,old-create-date=,current-digest=$image_digest,current-create-date=$image_created >> $ROLLING_LEDGER
            fi 
        fi 
    done
    IFS="$OLD_IFS"
}

clean_ledger() {
    check_for_ledger
    docker images --format '{{.Repository}}' 2>/dev/null | grep kasmweb > "$SCRIPT_DIR"/tmp_running_img.txt

    while IFS=',' read -r ledger_img other || [ -n "$ledger_img" ]; do
        if ! grep -qF "$ledger_img" "$SCRIPT_DIR"/tmp_running_img.txt; then
            sed -i "\|$ledger_img|d" "$ROLLING_LEDGER"
        fi
    done < "$ROLLING_LEDGER"

    rm -f "$SCRIPT_DIR"/tmp_running_img.txt
}


rollback=false
rollback_tag=""
latest=false
update=false
clean_up=false
date_image=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        -r|--rollback)
            rollback=true
            if [[ $# -ge 2 && ! "$2" =~ ^- ]];then
                rollback_tag="$2"
                echo "=== Rolling back to $rollback_tag ==="
                shift 2
            else
                echo "=== Rolling back one image release ==="
                shift
            fi
        ;;
        -r=*|--rollback=*)
            rollback=true
            rollback_tag="${1#*=}"
            echo "=== Rolling back to $rollback_tag ==="
            shift
        ;;
        -l|--latest)
            latest=true
            echo "=== Switching to using latest rolling image ==="
            shift
        ;;
        -u|--update)
            update=true
            shift
        ;;
        -h|--help)
            show_help
            exit 0
        ;;
        -c|--cleanup)
            clean_up=true
            echo "=== Removing old tracked images ==="
            shift
        ;;
        -d|--date)
            date_image=true
            shift
        ;;
        *)
            echo "Unknown options: $1"
            echo "Use -h or --help to see options"
            exit 1
        ;;
    esac
done

if [[ "$latest" == true ]]; then
    use_latest
    exit 0
fi

if  [[ "$update" == true ]]; then
    update_rolling
    exit 0
fi

if [[ "$clean_up" == true ]]; then
    clean_ledger
    exit 0
fi

# After a rollback, the tracking ledger will be updated automatically upon a restart 
# or you can manually invoke the this script again with the "-u" flag
if [[ "$rollback" == true ]]; then
    update_rolling
    if [[ -n "$rollback_tag" ]]; then
        use_specific_tag "$rollback_tag"
    else
        rollback_image_tag
    fi
    echo "Restart Kasm for rollback to be completed"
    exit 0
fi

if [[ "$date_image" == true ]]; then
    show_image_date
    exit 0
fi