#!/bin/bash

DOCKER=`which docker.io`

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

CURL=`which curl`

COMPUTE_COUNTER=0

ARVADOS_DOMAIN=dev.arvados

function usage {
    echo >&2
    echo >&2 "usage: $0 (start|stop|restart|reset|test) [options]"
    echo >&2
    echo >&2 "start    run new or restart stopped arvados containers"
    echo >&2 "stop     stop arvados containers"
    echo >&2 "restart  stop and then start arvados containers"
    echo >&2 "reset    stop and delete containers WARNING: this will delete the data inside Arvados!"
    echo >&2 "test     run tests"
    echo >&2
    echo >&2 "$0 options:"
    echo >&2 "  -b[bridge], --bridge[=bridge] Docker bridge (default bridge docker0)"
    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 "  -p, --keepproxy               Keepproxy server"
    echo >&2 "  -h, --help                    Display this help and exit"
    echo >&2 "      --domain=dns.domain       DNS domain used by containers (default dev.arvados)"
    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 bridge_ip_address {
    local bridge_name=$1
    # FIXME: add a more robust check here.
    # because ip command could be mising, multiple docker bridges could be there.. etc.
    echo $(ip --oneline --family inet addr show dev "$bridge_name" | awk '{ print $4 }'| cut -d/ -f1 )
}

function start_container {
    bridge_ip=$(bridge_ip_address "$bridge")

    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=$bridge_ip --dns-search=compute.$ARVADOS_DOMAIN --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=$bridge_ip --dns-search=compute.$ARVADOS_DOMAIN --hostname compute$COMPUTE_COUNTER -P --privileged --name $name"
        let COMPUTE_COUNTER=$(($COMPUTE_COUNTER + 1))
      else
        args="$args --dns=$bridge_ip --dns-search=$ARVADOS_DOMAIN --hostname ${name#_server} --name $name"
      fi
    fi
    if [[ "$3" != '' ]]; then
      local volume="$3"
      args="$args --volumes-from $volume"
    fi
    if [[ "$4" != '' ]]; then
      local link="$4"
      args="$args --link $link"
    fi
    local image=$5

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

    echo "Starting container: $name"
    `$DOCKER ps --all |grep -E "\b$name\b" -q`
    if [[ "$?" == "0" ]]; then
        echo "  $DOCKER start $name"
        container=`$DOCKER start $name`
    else
        echo "  $DOCKER run $args $image"
        container=`$DOCKER run $args $image`
    fi

    if [ "$?" != "0" -o "$container" = "" ]; then
      echo "Unable to start container"
      exit 1
    else
      echo "Started container: $container"
    fi

}

# Create a Docker data volume
function make_keep_volumes () {
    `$DOCKER ps --all |grep -E "\bkeep_data\b" -q`
    if [[ "$?" == "0" ]]; then
      return
    fi
    docker create -v /keep-data --name keep_data arvados/keep
}

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
    local start_keepproxy=false
    local bridge="docker0"
    local

    # NOTE: This requires GNU getopt (part of the util-linux package on Debian-based distros).
    local TEMP=`getopt -o d::s::b:a::cw::nkpvh \
                  --long doc::,sso::,api::,bridge:,compute,workbench::,nameserver,keep,keepproxy,vm,help,domain:: \
                  -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
            -b | --bridge)
                case "$2" in
                    *)  start_bridge=$2; shift 2 ;;
                esac
                ;;
            -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
                ;;
            -p | --keepproxy)
                start_keepproxy=true
                shift
                ;;
            --domain)
                case "$2" in
                    *) ARVADOS_DOMAIN="$2"; shift 2 ;;
                esac
                ;;
            --)
                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 &&
          $start_keepproxy == 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
        start_keepproxy=true
    fi

    if [[ $start_nameserver != false ]]
    then
      $DOCKER ps | grep skydns >/dev/null
      need_skydns="$?"

      $DOCKER ps | grep skydock >/dev/null
      need_skydock="$?"

      if [[ "$need_skydns" != 0 || "$need_skydock" != 0 ]]
      then
          # skydns and skydock need to both be running before everything else.
          # If they are not running we need to shut everything down and start
          # over, otherwise DNS will be broken and the containers won't find each other.
          do_stop
          need_skydns=1
          need_skydock=1
      fi

      # We rely on skydock and skydns for dns discovery between the slurm controller and compute nodes,
      # so make sure they are running
      $DOCKER ps | grep skydns >/dev/null
      if [[ $need_skydns != "0" ]]; then
        echo "Detecting bridge '$bridge' IP for crosbymichael/skydns"
        bridge_ip=$(bridge_ip_address "$bridge")

        echo "Starting crosbymichael/skydns container..."
        $DOCKER rm "skydns" 2>/dev/null
        echo $DOCKER run -d -p $bridge_ip:53:53/udp --name skydns crosbymichael/skydns -nameserver 8.8.8.8:53 -domain arvados
        $DOCKER run -d -p $bridge_ip:53:53/udp --name skydns crosbymichael/skydns -nameserver 8.8.8.8:53 -domain arvados
      fi
      $DOCKER ps | grep skydock >/dev/null
      if [[ "$need_skydock" != "0" ]]; then
        echo "Starting crosbymichael/skydock container..."
        $DOCKER rm "skydock" 2>/dev/null
        echo $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
        $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_sso != false ]]
    then
        start_container "$start_sso:443" "sso_server" '' '' "arvados/sso"
    fi

    if [[ $start_api != false ]]
    then
      if [[ $start_sso != false ]]; then
        start_container "$start_api:443" "api_server" '' "sso_server:sso" "arvados/api"
      else
        start_container "$start_api:443" "api_server" '' '' "arvados/api"
      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
        start_container "25107:25107" "keep_server_0" \
            "keep_data" \
            "api_server:api" \
            "arvados/keep"
        start_container "25108:25107" "keep_server_1" \
            "keep_data" \
            "api_server:api" \
            "arvados/keep"
    fi

    if [[ $start_keepproxy != false ]]
    then
        start_container "9902:9100" "keepproxy_server" '' \
            "api_server:api" \
            "arvados/keepproxy"
    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 "" "workbench_server" '' "" "arvados/workbench"
    fi

    if [[ $start_api != false ]]
    then
        if [[ -f "api/generated/superuser_token" ]]
        then
          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
        fi
    fi

    if [ "$(awk '($1 == "nameserver"){print $2; exit}' </etc/resolv.conf)" != "$bridge_ip" ]; then
        echo
        echo "******************************************************************"
        echo "To access Arvados you must add the Arvados nameserver to the top"
        echo "of your DNS configuration in /etc/resolv.conf:"
        echo "nameserver $bridge_ip"
        echo
        echo "Then run '$0 start' again"
        echo "******************************************************************"
        echo
    else
        while ! $CURL -k -L -f http://workbench.$ARVADOS_DOMAIN >/dev/null 2>/dev/null ; do
            echo "Waiting for Arvados to be ready."
            sleep 1
        done

        `$DOCKER ps |grep -E "\bdoc_server\b" -q`
        if [[ "$?" == "0" ]]; then
            echo
            echo "******************************************************************"
            echo "You can access the Arvados documentation at http://doc.$ARVADOS_DOMAIN"
            echo "******************************************************************"
            echo
        fi

        `$DOCKER ps |grep -E "\bworkbench_server\b" -q`
        if [[ "$?" == "0" ]]; then
            echo
            echo "********************************************************************"
            echo "You can access the Arvados workbench at http://workbench.$ARVADOS_DOMAIN"
            echo "********************************************************************"
            echo
        fi
    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=""
    local stop_keepproxy=""

    # NOTE: This requires GNU getopt (part of the util-linux package on Debian-based distros).
    local TEMP=`getopt -o dsacwnkpvh \
                  --long doc,sso,api,compute,workbench,nameserver,keep,keepproxy,vm,help,domain:: \
                  -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 -E "\bcompute[0-9]+\b" |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 ;;
            -p | --keepproxy )
                stop_keep="keepproxy_server" ; shift ;;
            --domain)
                case "$2" in
                    *) ARVADOS_DOMAIN="$2"; shift 2 ;;
                esac
                ;;
            --)
                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 == "" &&
          $stop_keepproxy == "" ]]
    then
        stop_doc=doc_server
        stop_sso=sso_server
        stop_api=api_server
        stop_compute=`$DOCKER ps |grep -E "\bcompute[0-9]+\b" |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"
        stop_keepproxy="keepproxy_server"
    fi

    $DOCKER stop $stop_doc $stop_sso $stop_api $stop_compute $stop_workbench $stop_nameserver $stop_keep $stop_keepproxy $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
}

function do_reset {
    for name in skydock skydns workbench_server shell doc_server keepproxy_server keep_server_0 keep_server_1 compute0 compute1 api_server keepproxy keep_data sso_server
    do
        `$DOCKER ps |grep -E "\b$name\b" -q`
        if [[ "$?" == "0" ]]; then
            echo "  $DOCKER stop $name"
            $DOCKER stop $name
        fi
        `$DOCKER ps --all |grep -E "\b$name\b" -q`
        if [[ "$?" == "0" ]]; then
            echo "  $DOCKER rm $name"
            $DOCKER rm $name
        fi
    done
}

if [ "$DOCKER" == '' ]
then
  echo "Docker not found. Please install it first."
  exit 2
fi

if [ "$CURL" == '' ]
then
  echo "Curl not found. Please install it first."
  exit 3
fi

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 $@
        ;;
    reset)
        shift
        do_reset $@
        ;;
    *)
        usage
        exit 1
        ;;
esac