#!/bin/bash
#
# kubectx(1) is a utility to manage and switch between kubectl contexts.

# Copyright 2017 Google Inc.
#
# 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.

[[ -n $DEBUG ]] && set -x

set -eou pipefail
IFS=$'\n\t'

KUBECTX="${HOME}/.kube/kubectx"

usage() {
  cat <<"EOF"
USAGE:
  kubectx                   : list the contexts
  kubectx <NAME>            : switch to context <NAME>
  kubectx -                 : switch to the previous context
  kubectx <NEW_NAME>=<NAME> : rename context <NAME> to <NEW_NAME>
  kubectx <NEW_NAME>=.      : rename current-context to <NEW_NAME>
  kubectx -d <NAME>         : delete context <NAME> ('.' for current-context)
                              (this command won't delete the user/cluster entry
                              that is used by the context)

  kubectx -h,--help         : show this message
EOF
  exit 1
}

current_context() {
  kubectl config view -o=jsonpath='{.current-context}'
}

get_contexts() {
  kubectl config get-contexts -o=name | sort -n
}

list_contexts() {
  set -u pipefail
  local cur
  cur="$(current_context)"

  local yellow darkbg normal
  yellow=$(tput setaf 3)
  darkbg=$(tput setab 0)
  normal=$(tput sgr0)

  for c in $(get_contexts); do
  if [[ "${c}" = "${cur}" ]]; then
    echo "${darkbg}${yellow}${c}${normal}"
  else
    echo "${c}"
  fi
  done
}

read_context() {
  if [[ -f "${KUBECTX}" ]]; then
    cat "${KUBECTX}"
  fi
}

save_context() {
  local saved
  saved="$(read_context)"

  if [[ "${saved}" != "${1}" ]]; then
    printf %s "${1}" > "${KUBECTX}"
  fi
}

switch_context() {
  kubectl config use-context "${1}"
}

set_context() {
  local prev
  prev="$(current_context)"

  switch_context "${1}"

  if [[ "${prev}" != "${1}" ]]; then
    save_context "${prev}"
  fi
}

swap_context() {
  local ctx
  ctx="$(read_context)"
  if [[ -z "${ctx}" ]]; then
    echo "error: No previous context found." >&2
    exit 1
  fi
  set_context "${ctx}"
}

user_of_context() {
  # TODO(ahmetb) no longer used, consider deleting
  kubectl config view \
    -o=jsonpath="{.contexts[?(@.name==\"${1}\")].context.user}"
}

cluster_of_context() {
  # TODO(ahmetb) no longer used, consider deleting
  kubectl config view \
    -o=jsonpath="{.contexts[?(@.name==\"${1}\")].context.cluster}"
}

context_exists() {
  grep -q ^"${1}"\$ <(kubectl config get-contexts -o=name)
}

rename_context() {
  local old_name="${1}"
  local new_name="${2}"

  if [[ "${old_name}" == "." ]]; then
    old_name="$(current_context)"
  fi

  # TODO(ahmetb) old_user and old_cluster are no longer used, clean up
  local old_user old_cluster
  old_user="$(user_of_context "${old_name}")"
  old_cluster="$(cluster_of_context "${old_name}")"
  if [[ -z "$old_user" || -z "$old_cluster" ]]; then
    echo "error: Cannot retrieve context ${old_name}."  >&2
    exit 1
  fi

  if context_exists "${new_name}"; then
    echo "Context \"${new_name}\" exists, deleting..." >&2
    kubectl config delete-context "${new_name}" 1>/dev/null 2>&1
  fi

  kubectl config rename-context "${old_name}" "${new_name}"
}

delete_context() {
  local ctx
  ctx="${1}"
  if [[ "${ctx}" == "." ]]; then
    ctx="$(current_context)"
  fi
  echo "Deleting context \"${ctx}\"..." >&2
  kubectl config delete-context "${ctx}"
}

main() {
  if [[ "$#" -eq 0 ]]; then
    list_contexts
  elif [[ "${1}" == "-d" ]]; then
    if [[ "$#" -lt 2 ]]; then
      echo "error: missing context NAME" >&2
      usage
    elif [[ "$#" -gt 2 ]]; then
      echo "error: too many arguments" >&2
      usage
    fi
    delete_context "${2}"
  elif [[ "$#" -gt 1 ]]; then
    echo "error: too many arguments" >&2
    usage
  elif [[ "$#" -eq 1 ]]; then
    if [[ "${1}" == "-" ]]; then
      swap_context
    elif [[ "${1}" == '-h' || "${1}" == '--help' ]]; then
      usage
    elif [[ "${1}" =~ ^-(.*) ]]; then
      echo "error: unrecognized flag \"${1}\"" >&2
      usage
    elif [[ "${1}" =~ (.+)=(.+) ]]; then
      rename_context "${BASH_REMATCH[2]}" "${BASH_REMATCH[1]}"
    else
      set_context "${1}"
    fi
  else
    usage
  fi
}

main "$@"
