#!/bin/bash -e

show_help() {
    echo "usage: check-state <jq-filter> <comparison> <expected-res> [error-message]"
    echo "       print-state <jq-filter>"
    echo "       change-snap-channel <snap-name> <channel>"
    echo "       force-autorefresh"
    echo "       prevent-autorefresh"
    echo "       wait-for-autorefresh [last-change-id]"
    echo "       wait-for-snap-autorefresh <snap-name> [last-change-id]"
    echo ""
    echo "The tool is used to manage the snapd state by editing the state.json file."
    echo "Also it is used to check state and refresh status"
}

print_state() {
    JQ_FILTER=$1
    if [ -z "$JQ_FILTER" ]; then
        echo "snapd-state: jq-filter is a required parameter"
        exit 1
    fi
    gojq -r "$JQ_FILTER" < /var/lib/snapd/state.json
}

check_state() {
    JQ_FILTER=$1
    COMP=$2
    RES=$3
    FAIL_MSG="${4:-failed values comparison}"

    if [ -z "$JQ_FILTER" ] || [ -z "$COMP" ] || [ -z "$RES" ]; then
        echo "snapd-state: jq-filter, comparison and expected-res are required parameters"
        exit 1
    fi

    if [ "$COMP" = "!=" ]; then
        if [ "$(print_state "$JQ_FILTER")" = "$RES" ]; then
            echo "snapd-state: $FAIL_MSG"
            exit 1
        fi
    elif [ "$COMP" = "=" ]; then
        if [ "$(print_state "$JQ_FILTER")" != "$RES" ]; then
            echo "snapd-state: $FAIL_MSG"
            exit 1
        fi
    else
        echo "snapd-state: invalid comparison operator: $COMP"
        exit 1
    fi
}

change_snap_channel() {
    local SNAP="$1"
    local CHANNEL="$2"
    if [ -z "$SNAP" ] || [ -z "$CHANNEL" ]; then
        echo "snapd-state: snap and channel are required parameters"
        exit 1
    fi
    gojq ".data.snaps[\"$SNAP\"].channel = \"$CHANNEL\"" < /var/lib/snapd/state.json > /var/lib/snapd/state.json.new
    mv /var/lib/snapd/state.json.new /var/lib/snapd/state.json
}

force_autorefresh() {
    gojq ".data[\"last-refresh\"] = \"2007-08-22T09:30:44.449455783+01:00\"" < /var/lib/snapd/state.json > /var/lib/snapd/state.json.new
    mv /var/lib/snapd/state.json.new /var/lib/snapd/state.json
}

prevent_autorefresh() {
    gojq ".data[\"last-refresh\"] = \"$(date +%Y-%m-%dT%H:%M:%S%:z)\"" < /var/lib/snapd/state.json > /var/lib/snapd/state.json.new
    mv /var/lib/snapd/state.json.new /var/lib/snapd/state.json
}

_wait_autorefresh(){
    local LAST_CHANGE_ID="$1"
    local RETRIES="$2"
    local CHECK_MESSAGE="$3"

    CHANGE_ID="$LAST_CHANGE_ID"
    for _ in $(seq "$RETRIES"); do
      # get last 2 lines of snap changes (the last one is always empty), match
      # auto-refresh change; only proceed if the change has greater change id
      # than the previously matched auto-refresh (this way we can match
      # consecutive auto-refreshes).
      if CHANGES=$(snap changes | tail -2 | grep "$CHECK_MESSAGE"); then
        CHANGE_ID=$(echo "$CHANGES" | awk '{print $1}')
        if [ "$CHANGE_ID" -gt "$LAST_CHANGE_ID" ]; then
          break
        fi
      fi
      snap debug ensure-state-soon
      sleep 1
    done
    if [ "$LAST_CHANGE_ID" -eq "$CHANGE_ID" ]; then
      echo "snapd-state: expected a new auto-refresh change with id greater than $LAST_CHANGE_ID, but it didn't happen"
      exit 1
    fi
    echo "$CHANGE_ID"
}

wait_for_autorefresh() {
    local LAST_CHANGE_ID="${1:-1}"

    _wait_autorefresh "$LAST_CHANGE_ID" 200 "Done.*Auto-refresh"
}

wait_for_snap_autorefresh() {
    local SNAP="$1"
    local LAST_CHANGE_ID="${2:-1}"

    if [ -z "$SNAP" ]; then
        echo "snapd-state: snap-name is a required parameter"
        exit 1
    fi

    _wait_autorefresh "$LAST_CHANGE_ID" 120 "Done.*Auto-refresh.*snap.*$SNAP"
}

main() {
    if [ $# -eq 0 ]; then
        show_help
        exit 0
    fi

    local subcommand="$1"
    local action=

    case "$1" in
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            action=$(echo "$subcommand" | tr '-' '_')
            shift
            ;;
    esac

    if [ -z "$(declare -f "$action")" ]; then
        echo "snapd-state: no such command: $subcommand"
        show_help
        exit 1
    fi

    "$action" "$@"
}

main "$@"
