#!/usr/bin/env bash
#
# Copyright (c) 2020 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

_mlops_targets="$(echo $(for i in $(find * -type d -name k8s);do name=${i##*tensorflow/}; name=${name%/training*}; echo $name; done))"
_mlops_targets_case=$(echo "$_mlops_targets"|sed 's/ /|/g')
_mlops_subcommands="deploy exec logs status undeploy"
_mlops_subcommands_case=$(echo "$_mlops_subcommands"|sed 's/ /|/g')
_mlops_subsubcommands="single-node multi-node"
_mlops_subsubcommands_case=$(echo "$_mlops_subsubcommands"|sed 's/ /|/g')
_dry_run=false
TARGETS='-h --help --dry-run options completion install-completion '"$_mlops_targets"

_usage()
{
  echo 'usage ' >&2
  echo ' '$(basename $0)' '$(echo $TARGETS|sed 's/ /|/g') >&2
  exit 0
}

_mlops_completion()
{
  __mlops_completion()
  {
    local _cur _prev _opts
    COMPREPLY=()
    _cur="${COMP_WORDS[COMP_CWORD]}"
    _prev="${COMP_WORDS[COMP_CWORD-1]}"
    _opts="__TARGETS__"
    case "${_prev}" in
      completion|install-completion)
        return 0
        ;;
      $_mlops_targets_case)
        local _options="$_mlops_subcommands options"
        COMPREPLY=( $(compgen -W "${_options}" -- ${_cur}) )
        return 0
        ;;
      $_mlops_subcommands_case)
        local _options="$_mlops_subsubcommands options"
        COMPREPLY=( $(compgen -W "${_options}" -- ${_cur}) )
        return 0
        ;;
      $_mlops_subsubcommands_case)
        return 0
        ;;
      --dry-run)
        local _options="$_mlops_targets options"
        COMPREPLY=( $(compgen -W "${_options}" -- ${_cur}) )
        return 0
        ;;
      -h|--help)
        local base=${COMP_WORDS[COMP_CWORD-2]}
        case "${base}" in
          *)
            return 0
            ;;
        esac
        ;;
      *)
        if (( ${#COMP_WORDS[@]} > 2 )); then
          local _options=$(eval "${COMP_WORDS[@]:0:$((${#COMP_WORDS[@]}-1))} options")
          COMPREPLY=( $(compgen -W "${_options}" -- ${_cur}) )
          return 0
        fi
        ;;
    esac
    COMPREPLY=($(compgen -W "${_opts}" -- ${_cur}))  
    return 0
  }
  complete -F __mlops_completion mlops
}

_mlops_install_completion()
{
  if [[ -d /usr/local/etc/bash_completion.d ]]; then
    $0 completion > /usr/local/etc/bash_completion.d/$(basename $0)
  fi
}

_mlops_exec()
{
  local _container _pod _mpijob _container
  case "$1" in 
    multi-node)
      _mpijob=$(kubectl get pods -oname 2>/dev/null | grep launcher | cut -c5-)
      if [[ -n $_mpijob ]]; then
        _mlops_invoke kubectl exec $_mpijob -c mpi-launcher -it -- bash
      fi
      ;;
    single-node)
      _job=$(kubectl get pods -oname 2>/dev/null | grep training | cut -c5-)
      if [[ -n $_job ]]; then
        _mlops_invoke kubectl exec $_job -c single-node -it -- bash
      fi
      ;;
    *)
      echo 'invalid option '$1 >&2
      exit 1
      ;;
  esac
}

_mlops_logs()
{
  local _mpijob
  case "$1" in 
    multi-node)
      _mpijob=$(kubectl get pods -oname 2>/dev/null | grep launcher | cut -c5-)
      if [[ -n $_mpijob ]]; then
        _mlops_invoke kubectl logs -f $_mpijob -c mpi-launcher
      fi
      ;;
    single-node)
      _job=$(kubectl get pods -oname 2>/dev/null | grep training | cut -c5-)
      if [[ -n $_job ]]; then
        _mlops_invoke kubectl logs -f $_job 
      fi
      ;;
    *)
      echo 'invalid option '$1 >&2
      exit 1
      ;;
  esac
}

_mlops_status()
{
  case "$1" in 
    multi-node)
      _count=$(kubectl get mpijobs.kubeflow.org 2>/dev/null|wc -l|sed -e 's/[[:space:]]*$//')
      if (( $_count > 0 )); then
        echo '📬 deployed'
      else
        echo '📭 not deployed'
      fi
      ;;
    single-node)
      kubectl get pods 2>/dev/null | grep 'resnet50v1-5-fp32-training' 1>/dev/null 2>/dev/null
      if (( $? == 0 )); then
        echo '📬 deployed'
      else
        echo '📭 not deployed'
      fi
      ;;
    *)
      echo 'invalid option '$1 >&2
      exit 1
      ;;
  esac
}

_mlops_kubectl_action()
{
  local _path=$1 _action=$2
  case "$_action" in
    apply)
      _mlops_invoke 'kubectl kustomize '$_path' | '\
        $'sed \'s/runAsUser:.*"\([0-9]*\)"/runAsUser: \\1/g\' | '\
        $'sed \'s/runAsGroup:.*"\([0-9]*\)"/runAsGroup: \\1/g\' | '\
        $'sed \'s/fsGroup:.*"\([0-9]*\)"/fsGroup: \\1/g\' | '\
        'kubectl apply -f -'
      ;;
    *)
      _mlops_invoke kubectl -k $_path $_action
      ;;
  esac
}

_mlops_invoke()
{
  if [[ $_dry_run == true ]]; then
    echo $@
  else
    eval $@
  fi
}

_main()
{
  local _kustomize_path=''
  while [[ "$#" -gt "0" && $1 =~ ^- ]]; do
    case "$1" in
      -h|--help)
          _usage
          exit 0
          ;;
      --dry-run)
          shift
          _dry_run=true
          ;;
      *)
          echo 'unknown option '$1
          shift
          exit 1
          ;;
    esac
  done
  case "$#" in
    0) 
      _usage
      ;;
    1)
      case "$1" in 
        completion)
          declare -a _targets=( $TARGETS )
          type _mlops_completion | sed '1,3d;$d' | eval sed 's/__TARGETS__/'"'${_targets[@]}'"'/g' | sed 's^\$_mlops_targets_case^'$_mlops_targets_case'^g' | sed 's^\$_mlops_targets^'"$_mlops_targets"'^g' | sed 's^\$_mlops_subcommands_case^'$_mlops_subcommands_case'^g' | sed 's^\$_mlops_subcommands^'"$_mlops_subcommands"'^g' | sed 's^\$_mlops_subsubcommands_case^'$_mlops_subsubcommands_case'^g' | sed 's^\$_mlops_subsubcommands^'"$_mlops_subsubcommands"'^g'
          ;;
        install-completion)
          _mlops_install_completion
          ;;
        options)
          echo $TARGETS' options '
          return 0
          ;;
        $_mlops_targets_case)
          echo 'missing subcommand for '$1 >&2
          _usage
          exit 1
          ;;
        *)
          echo 'unknown option '$1 >&2
          _usage
          exit 1
          ;;
      esac
      ;;
    2)
      case "$2" in 
        status)
          echo 'missing single-node|multi-node option ' >&2
          exit 1
          ;;
        exec)
          echo 'missing single-node|multi-node option ' >&2
          exit 1
          ;;
        logs)
          echo 'missing single-node|multi-node option ' >&2
          exit 1
          ;;
        options)
          echo 'deploy exec logs options status undeploy'
          return 0
          ;;
        deploy|undeploy)
          echo 'missing single-node|multi-node option ' >&2
          exit 1
          ;;
        *)
          echo 'unknown option '$2 >&2
          _usage
          exit 1
          ;;
      esac
      ;;
    3)
      _kustomize_path=$(find * -type d -name k8s|grep -v common)/mlops
      case "$3" in 
        single-node|multi-node)
          case "$2" in 
            deploy)
              _mlops_kubectl_action $_kustomize_path/$3 apply
              ;;
            undeploy)
              _mlops_kubectl_action $_kustomize_path/$3 delete
              ;;
            status)
              _mlops_status $3
              ;;
            exec)
              _mlops_exec $3
              ;;
            logs)
              _mlops_logs $3
              ;;
            options)
              echo '-h --help '$(echo $_mlops_subsubcommands)' options '
              return 0
              ;;
            *)
              echo 'unknown option '$2 >&2
              _usage
              ;;
          esac
          ;;
        *) 
          shift
          echo 'unknown option '$3 >&2
          ;;
      esac
      ;;
    *) 
      shift
      echo 'unknown options '$@ >&2
      ;;
  esac
}

_main $@
        

