#!/bin/bash

VNLK=/usr/bin/vnlk
VAR=/var/vnlk/update
VAE_LOG=$VAR/vae
SETTINGS=$VAR/settings.json
CENTOS=0
UBUNTU=0
REGISTRY_HOST=''
REGISTRY_USERNAME=''
REGISTRY_PASSWORD=''
MAXWAIT=600
PIDFILE=$VAR/update-vae.pid
DISCOVERY=/usr/bin/vae-discovery
VAE_IGNORE_LIST="vae-anpr vae-anpr-sdk"
CLEANUP_FILE=$VAR/.cleanup
CLEANUP_TIMEOUT=3600
VAE_LIST=''
VAE_PROJECT_ID=''
VAE_DEFAULT_PROJECT_ID=overcast
VAE_STAT=/var/vnlk/log/vae.stat

if [ -f $VAR/.disabled ]; then
    logger -p local0.notice -t update-vae "Updates disabled"
    exit 0
fi

function write-pid {
    echo $$ > $PIDFILE
}

function read-pid {
    if [ ! -f $PIDFILE ]; then
        return
    fi
    echo $(cat $PIDFILE)
}

function check-pid {
    local oldpid="$(read-pid)"
    if [ ! -z "$oldpid" ]; then
        local proc=$(ps -p $oldpid -o comm=)
        if [ "$proc" == "update-vae" ]; then
            die "Another instance of vae-update already running"
        fi
    fi
    write-pid
}

function log {
    local level=${2-info}
    echo >&2 -e "[$(date +"%Y-%m-%d %H:%M:%S")] ${1-}"
    logger -p local0.${level} -t update-vae "${1-}"
}

function logerr {
    log "${1-}" error
}

function die {
    local msg=$1
    local code=${2-1}
    log "$msg" error
    exit "$code"
}

function vae-log {
    local level=${2-info}
    echo -e "[$(date +"%Y-%m-%d %H:%M:%S")] ${1-}" >> "${VAE_LOG}/${CUR_VAE}.log"
    log "${1-}" $level
}

function vae-status {
    echo "${1-}" > "${VAE_LOG}/${CUR_VAE}.status"
}

function vae-version {
    echo "${1-}" > "${VAE_LOG}/${CUR_VAE}.version"
}

function vae-err {
    vae-log "${1-}" error
    vae-status error
}

function vae-ok {
    vae-log "${1-}"
    vae-status ok
}

function write-stat {
    local vae_list_str="$1"
    local vae_list=(${vae_list_str//,/ })
    local py_vae_list=''
    for img in ${vae_list[@]}; do
        py_vae_list="$py_vae_list,\"$img\""
    done
    py_vae_list="${py_vae_list:1}"
    local py="python"
    if type python3 >/dev/null 2>&1
    then
        py="python3"
    fi
    $py << END
import json, sys

vae_list = [$py_vae_list]
vae_updates_js = {}

def readFile(fn):
    content = ''
    try:
        with open(fn) as f:
            for line in f:
                content += line
        content = content.strip()
    except:
        print("could not read file {}: {}".format(fn, sys.exc_info()[1]))
    return content

for vae in vae_list:
    content = readFile("$VAE_LOG/{}.log".format(vae))
    status = readFile("$VAE_LOG/{}.status".format(vae))
    version = readFile("$VAE_LOG/{}.version".format(vae))
    vae_updates_js[vae]={"log":content, "status":status, "version": version}

with open("$VAE_STAT", "r") as f:
    lines = f.readlines()
with open("$VAE_STAT", "w") as f:
    for line in lines:
        if not line.startswith("STAT_VAE_UPDATES"):
            f.write(line)
    f.write("STAT_VAE_UPDATES="+json.dumps(vae_updates_js)+"\n")

END
}

function prepare {
    # Check Update flag first
    if grep -q 'VAE_IMAGE_UPDATE = no' /etc/vnlk/vnlk.conf; then
        log "VAE update is disabled. Exiting"
        exit 0
    fi
    # Get OS
    if grep -q "Ubuntu" /etc/os-release; then
        UBUNTU=1
    elif grep -q "CentOS" /etc/os-release; then
        CENTOS=1
    fi
    # Do nothing if we aren't running on neither Centos nor Ubuntu
    if [[ $CENTOS -eq 0 && $UBUNTU -eq 0 ]]; then exit 0; fi
    # Check for concurrent runs
    check-pid
    # Prepare directory for vae logs
    mkdir -p $VAE_LOG
    chown -R apl:apl $VAE_LOG

    log "Starting VAE update"
}

function get-updates-conf {
    if [ -f "$SETTINGS" ]; then
        local settings="$(cat $SETTINGS)"
    else
        die "Cannot get link update configuration"
    fi

    # Get registry URL
    REGISTRY_HOST=$(echo $settings | jq -r .list?.REGISTRY_HOST?)
    if [ -z "$REGISTRY_HOST" ] || [ "$REGISTRY_HOST" == "null" ]; then
        die "Cannot get registry host" 2
    fi
    REGISTRY_USERNAME=$(echo $settings | jq -r .list?.REGISTRY_USERNAME?)
    if [ -z "$REGISTRY_USERNAME" ] || [ "$REGISTRY_USERNAME" == "null" ]; then
        die "Cannot get registry username" 2
    fi
    REGISTRY_PASSWORD=$(echo $settings | jq -r .list?.REGISTRY_PASSWORD?)
    if [ -z "$REGISTRY_PASSWORD" ] || [ "$REGISTRY_PASSWORD" == "null" ]; then
        die "Cannot get registry password" 2
    fi

    # GET VAE list as csv-formatted string
    VAE_LIST=$(echo $settings | jq -r .list?.VAE_UPDATES_LIST?)
    if [ -z "$VAE_LIST" ] || [ "$VAE_LIST" == "null" ]; then
        log "VAE updates list is missing. Will try to get updates from the registry"
        VAE_LIST=''
    fi
    # Get Project ID to build repo URLs
    VAE_PROJECT_ID=$(echo $settings | jq -r .list?.VAE_PROJECT_ID?)
    if [ -z "$VAE_PROJECT_ID" ] || [ "$VAE_PROJECT_ID" == "null" ]; then
        log "Cannot get vae project id. Using default value"
        VAE_PROJECT_ID=$VAE_DEFAULT_PROJECT_ID
    fi
    VAE_PROJECT_ID=${VAE_PROJECT_ID-$VAE_DEFAULT_PROJECT_ID}


    REGISTRY_URL="https://$REGISTRY_USERNAME:$REGISTRY_PASSWORD@$REGISTRY_HOST"
}

function browse-vae-updates {
    VAE_LIST=''
    log "browsing VAE updates in the registry"
    for repo in $(curl -s "$REGISTRY_URL/v2/_catalog?n=1000" | jq -r '.repositories[]'); do
        local img=${repo##*/}
        local prefix=${repo%%/*}
        local skip="no"
        if [[ "$prefix" == "overcast" && $img =~ ^vae- ]]; then
            for ignore in $VAE_IGNORE_LIST; do
                if [ "$img" == "$ignore" ]; then
                    skip="yes"
                    break
                fi
            done
            if [ "$skip" == "yes" ]; then
                continue
            fi
            latest_tag=$(curl -s "$REGISTRY_URL/v2/$repo/tags/list" | jq -r '.tags[]' | grep -P '\d+\.\d+\.\d+\.\d+' | sort -V | tail -1)
            if [ -z "$latest_tag" ]; then
                logerr "could not list tags for $repo"
                continue
            fi
            VAE_LIST="${VAE_LIST}$img:$latest_tag,"
        fi
    done
    [ -n "$VAE_LIST" ] && VAE_LIST=${VAE_LIST::-1}
}

function update-vae {
    # Try to search for VAE updates in the registry
    if [ -z "$VAE_LIST" ]; then
        browse-vae-updates
    fi

    local run_discovery=''
    local vae_list=(${VAE_LIST//,/ })
    local vae_updates_list=''
    for img_ver in ${vae_list[@]}; do
        local img=${img_ver%%:*}
        local ver=${img_ver##*:}
        local repo=$VAE_PROJECT_ID/$img:$ver
        local target_repo=overcast/$img:latest
        [ -z "$img" ] && continue
        [ -z "$ver" ] && continue
        export CUR_VAE=$img
        local errfile=$VAE_LOG/$CUR_VAE.err
        rm -rf $VAE_LOG/$img.*
        vae-version $ver
        vae_updates_list="$vae_updates_list,$img"
        if [[ "$(docker images --format '{{.Repository}}:{{.Tag}}'| grep $img:$ver 2>/dev/null)" == "" ]]; then
            vae-log "Started updating $img to $ver"
            docker pull $REGISTRY_HOST/$repo >/dev/null 2>$errfile
            if [ $? == 0 ]; then
                docker tag $REGISTRY_HOST/$repo $target_repo
                vae-ok "Finished updating $img to $ver"
                run_discovery='yes'
            else
                vae-err "Failed to update $img to $ver: $(cat $errfile | tr '\n' ' ')"
            fi
        else
            vae-ok "$img already at $ver"
        fi
    done

    if [ "$run_discovery" == "yes" ]; then
        discovery
    fi

    write-stat ${vae_updates_list:1}
}

function cleanup {
    local curtime="$(date -u +%s)"
    if [ -f $CLEANUP_FILE ]; then
        local last_cleanup_time="$(date -u +%s -r $CLEANUP_FILE)"
        let diff=$curtime-$last_cleanup_time
        if [ $diff -lt $CLEANUP_TIMEOUT ]; then
            return
        fi
    fi
    touch $CLEANUP_FILE
    # Remove docker data
    docker system prune -f

    # Remove old tags
    local prev_repo=''
    local count=1
    for img in $(docker images --format "{{.Repository}}:{{.Tag}}" | grep overcast/vae | grep -v latest | sort -rV); do
        local tag=${img##*:}
        local repo=${img%%:*}
        if [ "$repo" == "$prev_repo" ]; then
            let count=$count+1
        else
            count=1
        fi
        prev_repo=$repo
        if [ $count -gt 1 ]; then
            docker rmi $repo:$tag
        fi
    done
}

function discovery {
    if [ -f $DISCOVERY ]; then
        log "Running VAE discovery"
        timeout -k 301 300 $DISCOVERY 2>&1
    fi
}

function init-vae {
    docker login -u $REGISTRY_USERNAME -p $REGISTRY_PASSWORD $REGISTRY_HOST
    if [ $? != 0 ]; then
        die "docker login failed" 3
    fi
}

function main {
    prepare

    get-updates-conf

    init-vae

    update-vae

    cleanup
}

main
