#!/bin/bash


# look for NOMAD_ADDR and NOMAD_TOKEN
[ -e $HOME/.config/nomad ]  &&  source $HOME/.config/nomad


# If not running interactively, don't setup autocomplete
if [ ! -z "$PS1" ]; then
  # nomad/consul autocompletes
  if [ "$ZSH_VERSION" = "" ]; then
    which  nomad >/dev/null  &&  complete            -C $(which  nomad)  nomad
    which consul >/dev/null  &&  complete            -C $(which consul) consul
  else
    # https://apple.stackexchange.com/questions/296477/
    ( which compdef 2>&1 |fgrep -q ' not found' )  &&  autoload -Uz compinit  &&  compinit

    which  nomad >/dev/null  &&  autoload -U +X bashcompinit && bashcompinit
    which  nomad >/dev/null  &&  complete -o nospace -C $(which  nomad)  nomad
    which consul >/dev/null  &&  complete -o nospace -C $(which consul) consul
  fi
fi


function nom-app() {
  # finds the webapp related to given job/CWD and opens it in browser
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)

  _nom-url

  URL=$(echo "$URL" |head -1)

  [ "$URL" = "" ]  &&  echo "URL not found - is service running? try:\n nomad status $JOB"  &&  return
  open "$URL"
}


function nom-ssh() {
  # simple way to pop in (ssh-like) to a given job
  # Usage: [job name, eg: x-thumb]  -OR-  no args will use CWD to determine job
  [ $# -ge 1 ]  &&   JOB=$1
  [ $# -lt 1 ]  &&   JOB=$(nom-job-from-cwd)
  [ $# -ge 2 ]  &&  TASK=$2    # for rarer TaskGroup case where 2+ Tasks spec-ed in same Job
  [ $# -lt 2 ]  &&  TASK=http

  ALLOC=$(nomad job status $JOB |egrep -m1 '\srun\s' |cut -f1 -d' ')
  echo "nomad alloc exec -i -t -task $TASK $ALLOC"

  if [ $# -ge 3 ]; then
    shift
    shift
    nomad alloc exec -i -t -task $TASK $ALLOC "$@"
  else
    nomad alloc exec -i -t -task $TASK $ALLOC \
      sh -c '([ -e /bin/zsh ] && zsh) || ([ -e /bin/bash ] && bash)  ||  ([ -e /bin/sh ] && sh)'
  fi
}


function nom-sshn() {
  # simple way to pop in (ssh-like) to a given job with 2+ allocations/containers
  local N=${1:?"Usage: [container/allocation number, starting with 1]"}

  local ALLOC=$(nomad job status $JOB |egrep '\srun\s' |head -n $N |tail -1 |cut -f1 -d' ')
  echo "nomad alloc exec -i -t $ALLOC"

  nomad alloc exec -i -t $ALLOC \
    sh -c '([ -e /bin/zsh ] && zsh) || ([ -e /bin/bash ] && bash)  ||  ([ -e /bin/sh ] && sh)'
}


function nom-cp() {
  # copies a laptop local file into running deploy (avoids full pipeline just to see changes)

  # first, see if this is vscode sync-rsync
  local VSCODE=
  [ "$#" -ge 4 ]  &&  ( echo "$@" |fgrep -q .vscode )  &&  VSCODE=1

  if [ $VSCODE ]; then
    # fish out file name from what VSCode 'sync-rsync' package sends us -- should be 2nd to last arg
    local FILE=$(echo "$@" |rev |tr -s ' ' |cut -f2 -d' ' |rev)
    # switch dirs to make aliases work
    local DIR=$(dirname "$FILE")
    cd "$DIR"
    local BRANCH=$(git rev-parse --abbrev-ref HEAD)
    local JOB=$(nom-job-from-cwd)
    local ALLOC=$(nom-job-to-alloc)
    local TASK=http
    cd -

  else
    local FILE=${1:?"Usage: [src file, locally qualified while 'cd'-ed inside a repo]"}
    local BRANCH=$(git rev-parse --abbrev-ref HEAD)
    local JOB=$(nom-job-from-cwd)
    local ALLOC=$(nom-job-to-alloc)
    [ $# -ge 2 ]  &&  TASK=$2    # for rarer TaskGroup case where 2+ Tasks spec-ed in sam Job
    [ $# -lt 2 ]  &&  TASK=http
  fi

  # now split the FILE name into two pieces -- 'the root of the git tree' and 'the rest'
  local DIR=$(dirname "$FILE")
  local TOP=$(git -C "$DIR" rev-parse --show-toplevel)
  local REST=$(echo "$FILE" | perl -pe "s|^$TOP||; s|^/+||;")


  for var in FILE DIR TOP REST BRANCH JOB ALLOC; do
    echo $var="${(P)var}"
  done
  echo

  # If called by VSCode and sync-rsync.sites.remotePath is not "", try to rsync
  if [ $VSCODE ] && [ "/$REST" != "${@: -1}" ]; then
    local MAIN=
    [ "$BRANCH" = "main"   ]  &&  MAIN=true
    [ "$BRANCH" = "master" ]  &&  MAIN=true

    local RSYNC=
    [   $MAIN ]  &&  RSYNC=true
    [ ! $MAIN ]  &&  [ "$RSYNC_BRANCHES" ]  &&  RSYNC=true

    [ $RSYNC ] && { set +e; set -x; rsync "$@"; RSYNC_EXIT=$?; set +x; set -e; }

    # this is a special exception project where we DONT want to ALSO copy file to nomad deploy
    [ $MAIN ]  &&  [ "$JOB" = "ia-petabox" ]  &&  [ ! $NOM_CP_PETABOX_MAIN ]  &&  exit 0
  fi


  if [ "$JOB" = ""  -o  "$ALLOC" = "" ]; then
    # no relevant job & alloc found - nothing to do
    echo 'has this branch run a full pipeline and deployed a Review App yet?'
    return
  fi

  # HinD updated nomad clusters w/ latest nomad seem to *not* get the stdin close properly
  # (and thus hang).  So timeout/kill after 2s :( tracey 2024/3 )
  set +e
  cat "$FILE" | ( set -x; set +e; nomad alloc exec -i -task $TASK "$ALLOC" sh -c "timeout 2 cat >| '$REST'" )
  NOMAD_EXIT=$?
  # timeout should come back with a 124 or 143 which we treat as success
  ([ $NOMAD_EXIT -eq 124 ] || [ $NOMAD_EXIT -eq 143 ]) && NOMAD_EXIT=0

  if [ -n "${RSYNC_EXIT}" ]; then
    if [ ${RSYNC_EXIT} -ne 0 ]; then
      echo "Warning: rsync failed with exit code $RSYNC_EXIT"
    else
      echo "Info: rsync was successful"
    fi
  fi

  if [ $NOMAD_EXIT -ne 0 ]; then
    echo "Warning: nomad sync failed with exit code $NOMAD_EXIT"
  else
    echo "Info: nomad sync was successful"
  fi

  if [ ${RSYNC_EXIT:-0} -ne 0 ] || [ $NOMAD_EXIT -ne 0 ]; then
    return 1
  fi
}


function nom-logs() {
  # simple way to view logs for a given job
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)
  # NOTE: the 2nd $JOB is useful for when a job has 2+ tasks (eg: `kv` or DB/redis, etc.)
  nomad alloc logs -f -job $JOB http
}


function nom-logs-err() {
  # simple way to view logs for a given job
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)
  nomad alloc logs -stderr -f -job $JOB http
}


function nom-status() {
  # prints detailed status for a repo's service and deployment
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)

  line
  echo "nomad status $JOB"
  line
        nomad status $JOB | grep --color=always -iE 'unhealthy|healthy|$'
  line
  echo 'nomad alloc status -stats $(nom-job-to-alloc '$JOB')'
  line
        nomad alloc status -stats $(nom-job-to-alloc $JOB) | grep --color=always -iE 'unhealthy|healthy|Job Version.*|Node Name.*|$'
  line
}


function nom-urls() {
  # Lists all current urls for the services deployed to current nomad cluster (eg: webapps)
  # Ideally, this is a faster single-shot call.  But to avoid requiring either `consul` addr
  # and ACL token _in addition_ to `nomad` - we'll just use `nomad` directly instead.
  #   consul catalog services -tags
  for JOB in $(curl -sH "X-Nomad-Token: ${NOMAD_TOKEN?}" ${NOMAD_ADDR?}/v1/jobs \
    | jq -r '.[] | select(.Type=="service") | "\(.Name)"')
  do
    _nom-url
    echo $URL
  done |sort
}


function _nom-url() {
  # logically private helper function
  URL=$(curl -sH "X-Nomad-Token: ${NOMAD_TOKEN?}" ${NOMAD_ADDR?}/v1/job/$JOB \
    | jq -r '.TaskGroups[0].Services[0].Tags' \
    | fgrep . |tr -d '", '
  )
}


function nom-resubmit() {
  # Retrieves current job spec from nomad cluster and resubmits it to nomad.
  # Useful for when a job has exceedded a setup timeout, is (nonideally) marked 'dead', etc.
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)

  # make able to work with any namespace, unless user is limiting namespaces
  local ARGS=(--namespace='*')
  [ "$NOMAD_NAMESPACE" != "" ] && ARGS=

  nomad inspect "$ARGS" ${JOB?} |perl -pe 's/"Stop": true,//' |tee .$JOB

  # in case we are trying to _move_ an active/OK deploy
  nomad stop "$ARGS" ${JOB?}
  sleep 5

  curl -XPOST -H "Content-Type: application/json" -H "X-Nomad-Token: $NOMAD_TOKEN" -d @.${JOB?} \
    $NOMAD_ADDR/v1/jobs

  rm -f .$JOB
}


function d() {
  # show docker running containers and local images
  [ "$#" = "0" ]  &&  clear -x

  local SUDO=
  [ $(uname) = "Linux" ]  &&  local SUDO=sudo
  [ ! -e /usr/bin/docker ]  &&  local docker=podman

  $SUDO $docker ps -a --format "table {{.ID}}\t{{.Image}}\t{{.Status}}\t{{.Names}}\t{{.State}}" | $SUDO cat >| $HOME/.dockps
  chmod 666 $HOME/.dockps
  for i in STATE running restarting created paused removing exited dead; do
    cat $HOME/.dockps |egrep "$i$" |perl -pe 's/'$i'$//'
  done
  rm -f $HOME/.dockps

  line
  $SUDO $docker images
}

function nom() {
  # quick way to get an overview of a nomad server when ssh-ed into it
  d
  line
  nomad server members
  line
  nomad status
  line
}


function nom-job-from-cwd() {
  # print the nomad job name based on the current project
  # parse out repo info, eg: 'ia-petabox' -- ensure clone-over-ssh or clone-over-https work
  local GURL TMP GROUP_PROJECT PROJECT BRANCH SLUG JOB
  GURL=$(git config --get remote.origin.url)
  [[ "$GURL" =~ https:// ]]  &&  TMP=$(echo "$GURL" |cut -f4- -d/)
  [[ "$GURL" =~ https:// ]]  ||  TMP=$(echo "$GURL" |rev |cut -f1 -d: |rev)
  GROUP_PROJECT=$(echo "$TMP" |perl -pe 's/\.git//' |tr A-Z a-z |tr / -)

  PROJECT=$(git rev-parse --absolute-git-dir |egrep --color -o '.*?.git' |rev |cut -f2 -d/ |rev)
  BRANCH=$(git rev-parse --abbrev-ref HEAD)
	SLUG=$(echo "$BRANCH" |tr '/_.' '-' |tr A-Z a-z)
  JOB=$GROUP_PROJECT
  [ "$SLUG" = "main" -o "$SLUG" = "master" -o "$SLUG" = "staging" -o "$SLUG" = "production" ]  ||  JOB="${JOB}-${SLUG}"

  echo $(echo "$JOB" |cut -b1-63)
}



function nom-image-from-cwd() {
  # print the registry image based on the current project
  # parse out repo info, eg: 'ia-petabox' -- ensure clone-over-ssh or clone-over-https work
  local GURL GROUP_PROJECT BRANCH SLUG JOB
  GURL=$(git config --get remote.origin.url)
  [[ "$GURL" =~ https:// ]]  &&  GROUP_PROJECT=$(echo "$GURL" |cut -f4- -d/)
  [[ "$GURL" =~ https:// ]]  ||  GROUP_PROJECT=$(echo "$GURL" |rev |cut -f1 -d: |rev)

  BRANCH=$(git rev-parse --abbrev-ref HEAD)
	SLUG=$(echo "$BRANCH" |tr '/_.' '-' |tr A-Z a-z)
  echo $(echo "registry.archive.org/$GROUP_PROJECT/$SLUG" |cut -b1-63)
}




function nom-job-to-alloc() {
  # prints alloc of a given job (when in high-availability and 2+ allocations, picks one at random)
  # Usage: [job name, eg: x-thumb]  -OR-  no args will use CWD to determine job
  [ $# -eq 1 ]  &&  JOB=$1
  [ $# -ne 1 ]  &&  JOB=$(nom-job-from-cwd)
  nomad job status $JOB |egrep -m1 '\srun\s' |cut -f1 -d' '
}


function line () {
  # horizontal line break
	perl -e 'print "_"x100; print "\n\n";'
}


function nom-tunnel() {
  # Sets up an ssh tunnel in the background to be able to talk to nomad cluster's consul.
  [ "$NOMAD_ADDR" = "" ]  &&  echo "Please set NOMAD_ADDR environment variable first"  &&  return
  local HOST=$(echo "$NOMAD_ADDR" | sed 's/:4646\/*$//' |sed 's/^https*:\/\///')
  ssh -fNA -L 8500:localhost:8500 $HOST
}


function web-logs-tail() {
  # admin script that can more easily "tail -f" the caddy (JSON) web logs
  (
    set -x
    tail -f /var/log/caddy/access.log | jq -r '"\(.request.host)\(.request.uri)\t\t\(.request.headers."User-Agent")"'
  )
}


function nom-deploys-check() {
  # admin script that more thoroughly finds and prints out any deploys with any kind of issues
  nomad system gc
  sleep 1.5
  clear -x
  echo 'Desired\tStatus\tJob'
	for node in $(nomad node status |grep -Fv Eligibility |cut -f1 -d' '); do
    (
      node_name=$(nomad node status -short "$node" | awk '/^Name/ {print $3}')
      echo "=== Node: $node_name ($node) ==="
      curl -sH "Content-Type: application/json" -H "X-Nomad-Token: $NOMAD_TOKEN" \
        $NOMAD_ADDR/v1/node/$node/allocations \
          |jq -r '.[] | [.DesiredStatus, .ClientStatus, .JobID] | @tsv'
    ) | grep -pEv '^run\trunning\t'
  done

  line

  nomad status --namespace '*' |grep -Fv running
}


function nom-deploys-check-only() {
  # admin script that more thoroughly finds and prints out any deploys with any kind of issues
  # this one doesnt do `nomad system gc` and has an alternate approach
  curl -sH "X-Nomad-Token: ${NOMAD_TOKEN?}" "${NOMAD_ADDR?}/v1/jobs?namespace=*" \
    |jq  . \
    |grep -E '"Name":|"Status":' \
    |grep -B1 Status \
    |grep -v '"running"' \
    |grep -B1 '"Status"' \
    |grep -v /periodic-1 \
    |grep -A1 '"Name"' \
    |tac \
    |perl -ne 'chop; print; print "\n" if m/"Name"/;' \
    |perl -pe 's/"Name"//g' \
    |cut -f2- -d: \
    |tr -d '":,' \
    |perl -pe 's/^\s+//; s/ +/ /g' \
    |tr ' ' '\t' \
    |sort
}


function nom-supervisor-status() {
  # admin script that shows the supervisord status of all jobs in each HinD cluster's VMs
  # requires `ssh` access to each VM that is running HinD
  for NODE in $(nomad node status -t '{{range .}}{{.Name}}{{"\n"}}{{end}}' |sort |grep -E .); do
    echo "=== $NODE ==="
    ssh $NODE 'sudo podman exec -it hind sh -c "supervisorctl status"'
    echo
  done
}

function nom-hello-world() {
  for NODE in $(nomad node status -t '{{range .}}{{.Name}}{{"\n"}}{{end}}' |sort |grep -E .); do
    echo "=== $NODE ==="
    ssh $NODE 'sudo podman exec -it hind sh -c "podman -r run --rm -it docker.io/hello-world"'
    echo
  done
}
