#!/bin/bash -e

garlicPrefix=@garlicPrefix@
sshHost=@sshHost@
PATH=@PATH@

usage() { echo "Usage: garlic [-RFwv] trebuchet" 1>&2; exit 1; }

msg() {
  >&2 echo "garlic: $@"
}

findClosure() {
  what=$1
  from=$2
  mexp=$(nix-store -qR "$from" | grep -E -- "$what")
  n=$(echo "$mexp" | awk 'BEGIN { count=0 } NF { count++ } END { print count }')
  if [ $n -eq 0 ]; then
    >&2 echo "$exp: $what not found"
    exit 1
  fi
  if [ $n -gt 1 ]; then
    >&2 echo "$exp: multiple $what found"
    exit 1
  fi
  echo "$mexp"
}

findExperiment() {
  grep -o -- "/nix/store/.*-experiment" "$1"
}

findOutputDir() {
  garlic_sandbox=$(nix show-config |\
    grep extra-sandbox-paths |\
    grep -o '/garlic=[^ ]*' || true)

  if [ -z "$garlic_sandbox" ]; then
    >&2 echo "Missing extra-sandbox-paths /garlic mountpoint"
    >&2 echo "Check the ~/.config/nix/nix.conf file"
    exit 1
  fi

  mountdir_rel=$(echo "$garlic_sandbox" | sed 's@^/garlic=@@g')
  mountdir=$(readlink -f "$mountdir_rel")

  echo "$mountdir/cache"
}

drvFromOutput() {
  nix-store -q --deriver $1
}

checkTrebuchet() {
  if [ ! -e "$trebuchet" ]; then
    >&2 echo "$trebuchet: not found"
    exit 1
  fi

  if [ ! -f "$trebuchet" ]; then
    >&2 echo "$trebuchet: not a file"
    exit 1
  fi

  # FIXME: We need a better way to determine a trebuchet
  if [ -z "$(grep "This trebuchet launches" $trebuchet)" ]; then
    >&2 echo "$trebuchet: not a trebuchet"
    exit 1
  fi

  return 0
}

checkExperiment() {
  if [ ! -e "$experiment" ]; then
    >&2 echo "$experiment: not found"
    exit 1
  fi

  if [ ! -f "$experiment" ]; then
    >&2 echo "$experiment: not a file"
    exit 1
  fi

  # FIXME: We need a better way to determine a experiment
  if [ -z "$(grep "This is an experiment" $experiment)" ]; then
    >&2 echo "$experiment: not an experiment"
    exit 1
  fi

  return 0
}

checkMountpoint() {
  if [ ! -e "$garlicPrefix/garlic.control" ]; then
    >&2 echo "error: missing $garlicPrefix/garlic.control"
    >&2 echo "Is the mountpoint enabled?"
    exit 1
  fi
}

status_line() {
  unithash="$1"
  name="$2"
  status="$3"

  red=$(tput -T ansi setaf 1)
  green=$(tput -T ansi setaf 2)
  yellow=$(tput -T ansi setaf 3)
  color_reset=$(tput -T ansi sgr0)

  case $status in
    ok)
      color_st="$green"
      ;;
    run*)
      color_st="$yellow"
      ;;
    exit*)
      color_st="$red"
      ;;
    *)
      color_st=""
      ;;
  esac

  if [ $verbose ]; then
    >&2 printf "garlic: %s  %s%-9s%s  %s\n" \
      "$unithash" "$color_st" "$status" "$color_reset" "$name"
  fi
}

do_fetch() {
  expName=$(basename $experiment)
  user=$(ssh -G "$sshHost" | awk '/^user /{print $2}')
  exp=$garlicPrefix/$user/out/$expName

  if [ ! -e "$exp" ]; then
    echo "missing experiment: $exp"
    exit 1
  fi

  unitNameList=$(grep '^/nix/store/.*-unit$' $experiment |\
    sort -u |\
    sed 's@^/nix/store/@@')

  cwd=$(pwd)

  repeat=1
  bad=0
  while [ 1 ]; do
    repeat=0
    test $verbose && msg "Checking units $(date --rfc-3339=seconds)..."

    declare -A unit_names

    for unit in $unitNameList; do

      unit_hash="${unit%-unit}"
      unit_status="?"
      unit_name="?"

      if [ -e "$exp/$unit" ]; then
        st_file="$exp/$unit/status"
        conf_json="$exp/$unit/garlic_config.json"
        done_file="$exp/$unit/done"

        if [ -z "${unit_names[$unit_hash]}" ]; then
          if [ -e "$conf_json" ]; then
            unit_names+=([$unit_hash]=$(jq -r .unitName "$conf_json"))
            unit_name="${unit_names[$unit_hash]}"
          fi
        else
          unit_name="${unit_names[$unit_hash]}"
        fi

        if [ -e "$st_file" ]; then
          unit_status=$(cat "$st_file")
        fi
      fi

      status_line "$unit_hash" "$unit_name" "$unit_status"

      if [ ! -e "$done_file" ]; then
        repeat=1
      elif [ "$unit_status" != "ok" ]; then
        bad=1
      fi
    done

    if [ $repeat -eq 0 ]; then
      break
    else
      test $verbose && msg ""
    fi

    if [ $waitResults -eq 1 ]; then
      #echo "waiting 3 seconds to try again"
      sleep 3
    else
      break
    fi
  done

  if [ $repeat -eq 1 ]; then
    #>&2 echo "$exp: execution incomplete"
    exit 1
  fi

  if [ $bad -eq 1 ]; then
    msg "Some units failed, aborting"
    exit 1
  fi

  cd "$cwd"

  test $verbose && msg "execution complete, fetching results"

  mkdir -p "$outputDir"

  rsync -vrt --copy-links \
    --include='*/*/garlic_config.json' \
    --include='*/*/std*.log' \
    --include='*/*/*/std*.log' \
    --include='*/*/*/.garlic' \
    --include='*/*/*/.garlic/**' \
    --exclude='*/*/*/*' \
    "$exp" "$outputDir"

  if [ ! $enableKeep ]; then
    nix-build -E "(with import ./default.nix; \
      garlic.pp.store { \
        trebuchet = (import \"$trebuchetDrv\" ); \
        experiment = (import \"$experimentDrv\"); \
      })"

    rm -rf "$outputDir/$expName"
  fi
}

do_run() {

  >&2 $trebuchet
}

do_delete() {
  expName=$(basename $experiment)
  rm -rf $outputDir/$expName
}

waitResults=1
verbose=
operation=
target=
enableRun=
enableFetch=
enableKeep=
enableDelete=

while getopts "vwRFKD" o; do
  case "${o}" in
    R) enableRun=1 ;;
    F) enableFetch=1 ;;
    K) enableKeep=1 ;;
    D) enableDelete=1 ;;
    w) waitResults=0 ;;
    v) verbose=1 ;;
    *) usage ;;
  esac
done
shift $((OPTIND-1))
trebuchet="$1"

if [ -z "$enableRun" -a -z "$enableFetch" -a -z "$enableDelete" ]; then
  >&2 echo "missing operation"
  usage
fi

if [ -z "$trebuchet" ]; then
  >&2 echo "missing experiment"
  usage
fi

checkMountpoint

checkTrebuchet $trebuchet

outputDir=$(findOutputDir)

experiment=$(findExperiment "$trebuchet")
checkExperiment $experiment

trebuchetDrv=$(drvFromOutput $trebuchet)
experimentDrv=$(drvFromOutput $experiment)

if [ $enableRun ]; then do_run; fi
if [ $enableFetch ]; then do_fetch; fi
if [ $enableDelete ]; then do_delete; fi