#!/bin/bash

ENABLE_SSH=false
DOCKER=`which docker.io`

if [[ "$DOCKER" == "" ]]; then
    DOCKER=`which docker`
fi

COMPUTE_COUNTER=0

function usage {
    echo >&2
    echo >&2 "usage: $0 (start|stop|restart|test) [options]"
    echo >&2
    echo >&2 "$0 start/stop/restart options:"
    echo >&2 "  -d[port], --doc[=port]        Documentation server (default port 9898)"
    echo >&2 "  -w[port], --workbench[=port]  Workbench server (default port 9899)"
    echo >&2 "  -s[port], --sso[=port]        SSO server (default port 9901)"
    echo >&2 "  -a[port], --api[=port]        API server (default port 9900)"
    echo >&2 "  -c, --compute                 Compute nodes (starts 2)"
    echo >&2 "  -v, --vm                      Shell server"
    echo >&2 "  -n, --nameserver              Nameserver"
    echo >&2 "  -k, --keep                    Keep servers"
    echo >&2 "  --ssh                         Enable SSH access to server containers"
    echo >&2 "  -h, --help                    Display this help and exit"
    echo >&2
    echo >&2 "  If no options are given, the action is applied to all servers."
    echo >&2
    echo >&2 "$0 test [testname] [testname] ..."
    echo >&2 "  By default, all tests are run."
}

function ip_address {
    local container=$1
    echo `$DOCKER inspect $container  |grep IPAddress |cut -f4 -d\"`
}

function start_container {
    local args="-d -i -t"
    if [[ "$1" != '' ]]; then
      local port="$1"
      args="$args -p $port"
    fi
    if [[ "$2" != '' ]]; then
      local name="$2"
      if [[ "$name" == "api_server" ]]; then
        args="$args --dns=172.17.42.1 --dns-search=compute.dev.arvados --hostname api -P --name $name"
      elif [[ "$name" == "compute" ]]; then
        name=$name$COMPUTE_COUNTER
        # We need --privileged because we run docker-inside-docker on the compute nodes
        args="$args --dns=172.17.42.1 --dns-search=compute.dev.arvados --hostname compute$COMPUTE_COUNTER -P --privileged --name $name"
        let COMPUTE_COUNTER=$(($COMPUTE_COUNTER + 1))
      else
        args="$args --dns=172.17.42.1 --dns-search=dev.arvados --hostname ${name#_server} --name $name"
      fi
    fi
    if [[ "$3" != '' ]]; then
      local volume="$3"
      args="$args -v $volume"
    fi
    if [[ "$4" != '' ]]; then
      local link="$4"
      args="$args --link $link"
    fi
    local image=$5

    if $ENABLE_SSH
    then
      args="$args -e ENABLE_SSH=$ENABLE_SSH"
    fi

    `$DOCKER ps |grep -P "$name[^/]" -q`
    if [[ "$?" == "0" ]]; then
      echo "You have a running container with name $name -- skipping."
      return
    fi

    # Remove any existing container by this name.
    $DOCKER rm "$name" 2>/dev/null

    echo "Starting container:"
    #echo "  $DOCKER run --dns=127.0.0.1 $args $image"
    echo "  $DOCKER run $args $image"
    container=`$DOCKER run $args $image`
    if [[ "$?" != "0" ]]; then
      echo "Unable to start container"
      exit 1
    fi
    if [[ "$name" == "compute" || $ENABLE_SSH ]];
    then
      ip=$(ip_address $container )
      echo
      echo "You can ssh into the container with:"
      echo
      echo "    ssh root@$ip"
      echo
    fi

    if [[ "$name" == "doc_server" ]]; then
      echo
      echo "*****************************************************************"
      echo "You can access the Arvados documentation at http://localhost:${port%:*}"
      echo "*****************************************************************"
      echo
    fi

    if [[ "$name" == "workbench_server" ]]; then
      echo
      echo "*****************************************************************"
      echo "You can access the Arvados workbench at http://localhost:${port%:*}"
      echo "*****************************************************************"
      echo
   fi


}

declare -a keep_volumes

# Initialize the global `keep_volumes' array. If any keep volumes
# already appear to exist (mounted volumes with a top-level "keep"
# directory), use them; create temporary volumes if necessary.
#
function make_keep_volumes () {
    # Mount a keep volume if we don't already have one
    for mountpoint in $(cut -d ' ' -f 2 /proc/mounts); do
      if [[ -d "$mountpoint/keep" && "$mountpoint" != "/" ]]; then
        keep_volumes+=($mountpoint)
      fi
    done

    # Create any keep volumes that do not yet exist.
    while [ ${#keep_volumes[*]} -lt 2 ]
    do
        new_keep=$(mktemp -d)
        echo >&2 "mounting 512M tmpfs keep volume in $new_keep"
        sudo mount -t tmpfs -o size=512M tmpfs $new_keep
        mkdir $new_keep/keep
        keep_volumes+=($new_keep)
    done
}

function do_start {
    local start_doc=false
    local start_sso=false
    local start_api=false
    local start_compute=false
    local start_workbench=false
    local start_vm=false
    local start_nameserver=false
    local start_keep=false

    # NOTE: This requires GNU getopt (part of the util-linux package on Debian-based distros).
    local TEMP=`getopt -o d::s::a::cw::nkvh \
                  --long doc::,sso::,api::,compute,workbench::,nameserver,keep,vm,help,ssh \
                  -n "$0" -- "$@"`

    if [ $? != 0 ] ; then echo "Use -h for help"; exit 1 ; fi

    # Note the quotes around `$TEMP': they are essential!
    eval set -- "$TEMP"

    while [ $# -ge 1 ]
    do
        case $1 in
            -d | --doc)
                case "$2" in
                    "") start_doc=9898; shift 2 ;;
                    *)  start_doc=$2; shift 2 ;;
                esac
                ;;
            -s | --sso)
                case "$2" in
                    "") start_sso=9901; shift 2 ;;
                    *)  start_sso=$2; shift 2 ;;
                esac
                ;;
            -a | --api)
                case "$2" in
                    "") start_api=9900; shift 2 ;;
                    *)  start_api=$2; shift 2 ;;
                esac
                ;;
            -c | --compute)
                start_compute=2
                shift
                ;;
            -w | --workbench)
                case "$2" in
                    "") start_workbench=9899; shift 2 ;;
                    *)  start_workbench=$2; shift 2 ;;
                esac
                ;;
            -v | --vm)
                start_vm=true
                shift
                ;;
            -n | --nameserver)
                start_nameserver=true
                shift
                ;;
            -k | --keep)
                start_keep=true
                shift
                ;;
            --ssh)
                # ENABLE_SSH is a global variable
                ENABLE_SSH=true
                shift
                ;;
            --)
                shift
                break
                ;;
            *)
                usage
                exit 1
                ;;
        esac
    done

    # If no options were selected, then start all servers.
    if [[ $start_doc == false &&
          $start_sso == false &&
          $start_api == false &&
          $start_compute == false &&
          $start_workbench == false &&
          $start_vm == false &&
          $start_nameserver == false &&
          $start_keep == false ]]
    then
        start_doc=9898
        start_sso=9901
        start_api=9900
        start_compute=2
        start_workbench=9899
        start_vm=true
        start_nameserver=true
        start_keep=true
    fi

    if [[ $start_sso != false ]]
    then
        start_container "$start_sso:443" "sso_server" '' '' "arvados/sso"
    fi

    if [[ $start_api != false ]]
    then
        start_container "$start_api:443" "api_server" '' "sso_server:sso" "arvados/api"
    fi

    if [[ $start_nameserver != false ]]
    then
      # We rely on skydock and skydns for dns discovery between the slurm controller and compute nodes,
      # so make sure they are running
      $DOCKER images | grep skydns >/dev/null
      if [[ "$?" != "0" ]]; then
        echo "Downloading crosbymichael/skydns..."
        $DOCKER pull crosbymichael/skydns
      fi
      $DOCKER ps | grep skydns >/dev/null
      if [[ "$?" != "0" ]]; then
        echo "Starting crosbymichael/skydns container..."
        $DOCKER rm "skydns" 2>/dev/null
        $DOCKER run -d -p 172.17.42.1:53:53/udp --name skydns crosbymichael/skydns -nameserver 8.8.8.8:53 -domain arvados
      fi
      $DOCKER images | grep skydock >/dev/null
      if [[ "$?" != "0" ]]; then
        echo "Downloading crosbymichael/skydock..."
        $DOCKER pull crosbymichael/skydock
      fi
      $DOCKER ps | grep skydock >/dev/null
      if [[ "$?" != "0" ]]; then
        echo "Starting crosbymichael/skydock container..."
        $DOCKER rm "skydock" 2>/dev/null
        $DOCKER run -d -v /var/run/docker.sock:/docker.sock --name skydock crosbymichael/skydock -ttl 30 -environment dev -s /docker.sock -domain arvados -name skydns
      fi
    fi

    if [[ $start_compute != false ]]
    then
        for i in `seq 0 $(($start_compute - 1))`; do
          start_container "" "compute" '' "api_server:api" "arvados/compute"
        done
    fi

    if [[ $start_keep != false ]]
    then
        # create `keep_volumes' array with a list of keep mount points
        # remove any stale metadata from those volumes before starting them
        make_keep_volumes
        for v in ${keep_volumes[*]}
        do
            [ -f $v/keep/.metadata.yml ] && sudo rm $v/keep/.metadata.yml
        done
        start_container "25107:25107" "keep_server_0" \
            "${keep_volumes[0]}:/dev/keep-0" \
            "api_server:api" \
            "arvados/keep"
        start_container "25108:25107" "keep_server_1" \
            "${keep_volumes[1]}:/dev/keep-0" \
            "api_server:api" \
            "arvados/keep"
    fi

    if [[ $start_doc != false ]]
    then
        start_container "$start_doc:80" "doc_server" '' '' "arvados/doc"
    fi

    if [[ $start_vm != false ]]
    then
        start_container "" "shell" '' "api_server:api" "arvados/shell"
    fi

    if [[ $start_workbench != false ]]
    then
        start_container "$start_workbench:80" "workbench_server" '' "api_server:api" "arvados/workbench"
    fi

    if [ -d $HOME/.config/arvados ] || mkdir -p $HOME/.config/arvados
    then
        cat >$HOME/.config/arvados/settings.conf <<EOF
ARVADOS_API_HOST=$(ip_address "api_server")
ARVADOS_API_HOST_INSECURE=yes
ARVADOS_API_TOKEN=$(cat api/generated/superuser_token)
EOF
    fi

}

function do_stop {
    local stop_doc=""
    local stop_sso=""
    local stop_api=""
    local stop_compute=""
    local stop_workbench=""
    local stop_nameserver=""
    local stop_vm=""
    local stop_keep=""

    # NOTE: This requires GNU getopt (part of the util-linux package on Debian-based distros).
    local TEMP=`getopt -o dsacwnkvh \
                  --long doc,sso,api,compute,workbench,nameserver,keep,vm,help \
                  -n "$0" -- "$@"`

    if [ $? != 0 ] ; then echo "Use -h for help"; exit 1 ; fi

    # Note the quotes around `$TEMP': they are essential!
    eval set -- "$TEMP"

    while [ $# -ge 1 ]
    do
        case $1 in
            -d | --doc)
                stop_doc=doc_server ; shift ;;
            -s | --sso)
                stop_sso=sso_server ; shift ;;
            -a | --api)
                stop_api=api_server ; shift ;;
            -c | --compute)
                stop_compute=`$DOCKER ps |grep -P "compute\d+" |grep -v api_server |cut -f1 -d ' '` ; shift ;;
            -w | --workbench)
                stop_workbench=workbench_server ; shift ;;
            -n | --nameserver )
                stop_nameserver="skydock skydns" ; shift ;;
            -v | --vm )
                stop_vm="shell" ; shift ;;
            -k | --keep )
                stop_keep="keep_server_0 keep_server_1" ; shift ;;
            --)
                shift
                break
                ;;
            *)
                usage
                exit 1
                ;;
        esac
    done

    # If no options were selected, then stop all servers.
    if [[ $stop_doc == "" &&
          $stop_sso == "" &&
          $stop_api == "" &&
          $stop_compute == "" &&
          $stop_workbench == "" &&
          $stop_vm == "" &&
          $stop_nameserver == "" &&
          $stop_keep == "" ]]
    then
        stop_doc=doc_server
        stop_sso=sso_server
        stop_api=api_server
        stop_compute=`$DOCKER ps |grep -P "compute\d+" |grep -v api_server |cut -f1 -d ' '`
        stop_workbench=workbench_server
        stop_vm=shell
        stop_nameserver="skydock skydns"
        stop_keep="keep_server_0 keep_server_1"
    fi

    $DOCKER stop $stop_doc $stop_sso $stop_api $stop_compute $stop_workbench $stop_nameserver $stop_keep $stop_vm \
        2>/dev/null
}

function do_test {
    local alltests
    if [ $# -lt 1 ]
    then
        alltests="python-sdk api"
    else
        alltests="$@"
    fi

    for testname in $alltests
    do
        echo "testing $testname..."
        case $testname in
            python-sdk)
                do_start --api --keep --sso
                export ARVADOS_API_HOST=$(ip_address "api_server")
                export ARVADOS_API_HOST_INSECURE=yes
                export ARVADOS_API_TOKEN=$(cat api/generated/superuser_token)
                python -m unittest discover ../sdk/python
                ;;
            api)
                $DOCKER run -t -i arvados/api \
                    /usr/src/arvados/services/api/script/rake_test.sh
                ;;
            *)
                echo >&2 "unknown test $testname"
                ;;
        esac
    done
}

if [ $# -lt 1 ]
then
  usage
  exit 1
fi

case $1 in
    start)
        shift
        do_start $@
        ;;
    stop)
        shift
        do_stop $@
        ;;
    restart)
        shift
        do_stop $@
        do_start $@
        ;;
    test)
        shift
        do_test $@
        ;;
    *)
        usage
        exit 1
        ;;
esac