#!/bin/bash

set -o errexit
set -o nounset
set -o pipefail

# TiDB version
TIDB_VERSION="v2.0.8"

# Default Docker registry
DOCKER_REGISTRY=uhub.ucloud.cn

# Default Docker repository for TiDB Enterprise Installer image
DOCKER_REPOSITORY=pingcap

ROOT_PATH=`pwd`

# Default host storage path
HOST_STORAGE_PATH="${ROOT_PATH}/tidb-enterprise-installer/data"
DOCKER_STORAGE_PATH="/ops-installer/worker/data"
STATE_FILE="${HOST_STORAGE_PATH}/state"
DASHBOARD_INFO_FILE="${HOST_STORAGE_PATH}/dashboard"
TIDB_INFO_FILE="${HOST_STORAGE_PATH}/tidb"

# Get from the client or assume a default location
HOST_DOCKER_HOST=${DOCKER_HOST:-/var/run/docker.sock}

# Output colour
CYAN='\033[0;36m'
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# Exit codes
GENERAL_ERROR_EXIT_CODE=1         # General errors. All errors other than the ones related to either argument or command
UNKNOWN_COMMAND_EXIT_CODE=2       # Unknown command
INVALID_ARGUMENT_EXIT_CODE=3      # Unknown argument or argument's value is not specified
NON_VALID_USER_EXIT_CODE=4        # When starting the installer with a user that is not a member of the docker group
PRECONDITION_NOT_MET_EXIT_CODE=5  # Pre-condition checks that error

# Quickly detect if the docker socket location exists or not
validateDockerSocket(){
  if [ ! -S "${HOST_DOCKER_HOST}" ]; then
    echo -e "${RED}TiDB Enterprise could not verify the Docker socket (${HOST_DOCKER_HOST})." \
      "\nTo resolve the issue, verify that the Docker daemon is running on this host and" \
      "that you are using the correct Docker socket to connect to the daemon.${NC}"
    exit $PRECONDITION_NOT_MET_EXIT_CODE
  fi
}

# Verify that the user is a member of the docker group.
validateRunningUser() {
  local uuid=$(id -ru)
  if [[ $uuid -eq 0 ]]; then
    return
  fi

  local docker_group=$(id -nG | grep -E '(^|\s)docker($|\s)')
  if [ -z "$docker_group" ]; then
    echo -e "${RED}The user is not a member of the docker group.${NC}"
    echo -e "${RED}To resolve the issue, add the user to the docker group or install as a different user.${NC}"
    exit $NON_VALID_USER_EXIT_CODE
  fi
}

createStateFile() {
  echo "SUCCESS" > ${STATE_FILE}
}

validateStateFileExist() {
  if [[ ! -e ${STATE_FILE} ]]; then
    echo -e "${RED}You need to run '${0##*/} install' command first.${NC}"
    exit $PRECONDITION_NOT_MET_EXIT_CODE
  fi
}

validateStateFileNotExist() {
  if [[ -e ${STATE_FILE} ]]; then
    echo -e "${RED}You've already run '${0##*/} install' command.${NC}"
    exit $PRECONDITION_NOT_MET_EXIT_CODE
  fi
}

createAndValidateHostStoragePath() {
  if [[ ! -e ${HOST_STORAGE_PATH} ]]; then
    mkdir -p ${HOST_STORAGE_PATH}
  fi

  if [[ ! -r ${HOST_STORAGE_PATH} ]]; then
    printf "${RED}%s${NC}\n" "Host storage path ${HOST_STORAGE_PATH} exists but doesn't have read permissions for user '${USER}'."
    printf "${RED}%s${NC}\n" "Please supply the correct permissions for the host storage path."
    exit $GENERAL_ERROR_EXIT_CODE
  fi

  if [[ ! -w ${HOST_STORAGE_PATH} ]]; then
    printf "${RED}%s${NC}\n" "Host storage path ${HOST_STORAGE_PATH} exists but doesn't have write permissions for user '${USER}'."
    printf "${RED}%s${NC}\n" "Please supply the correct permissions for the host storage path."
    exit $GENERAL_ERROR_EXIT_CODE
  fi
}

printInstallHeader() {
  echo -e "${CYAN}================================================================================================${NC}"
  echo -e "${CYAN}TiDB Enterprise ${TIDB_VERSION} Installer${NC}"

  echo -e "${CYAN}Start setting up a new TiDB Enterprise installation.${NC}"
  echo -e "${CYAN}================================================================================================${NC}"
}

printHeader() {
  local command="$1"

  echo -e "${CYAN}================================================================================================${NC}"
  echo -e "${CYAN}TiDB Enterprise ${TIDB_VERSION} Installer${NC}"

  echo -e "${CYAN}Start ${command} the TiDB cluster.${NC}"
  echo -e "${CYAN}================================================================================================${NC}"
}

printStartFooter() {
  echo -e "${CYAN}================================================================================================${NC}"
  echo -e "${CYAN}TiDB Enterprise ${TIDB_VERSION} installation completed successfully${NC}"

  if [[ -r ${DASHBOARD_INFO_FILE} ]]; then
    dashboard_info=`cat ${DASHBOARD_INFO_FILE}`
    echo -e "${CYAN}${dashboard_info}${NC}"
  fi

  if [[ -r ${TIDB_INFO_FILE} ]]; then
    echo -e "${CYAN}You could use the MySQL client to connect to the following TiDB servers directly:${NC}"
    tidb_info=`cat ${TIDB_INFO_FILE}`
    echo -e "${CYAN}${tidb_info}${NC}"
  fi

  echo -e "${CYAN}================================================================================================${NC}"
}

# Perform pre-condition checks on the host before starting the installation
verifyHostPreconditions() {
  # Check if we can connect to the docker socket
  validateDockerSocket

  # Check if the user is a member of the docker group
  validateRunningUser
}

runContainer() {
  local command="$1"

  docker -H "unix://${HOST_DOCKER_HOST}" run \
    -v ${HOST_STORAGE_PATH}:${DOCKER_STORAGE_PATH} \
    --name tidb-enterprise-installer-${TIDB_VERSION} \
    -it --rm ${DOCKER_REGISTRY}/${DOCKER_REPOSITORY}/tidb-enterprise-installer:${TIDB_VERSION} /ops-installer/ops-installer ${command}
}

startJob() {
  printHeader starting
  verifyHostPreconditions
  validateStateFileExist

  runContainer start

  printStartFooter
}

startChoice() {
  while true; do
      read -p "Do you wish to start the TiDB cluster now? (y/N)" result
      case $result in
          [Yy]* )
            startJob
            break
            ;;
          [Nn]* )
            exit
            ;;
          * ) echo "Please answer y or N.";;
      esac
  done
}

deployJob() {
  printHeader deploying
  verifyHostPreconditions
  createAndValidateHostStoragePath
  validateStateFileExist

  runContainer deploy

  createStateFile
}

stopJob() {
  printHeader stopping
  verifyHostPreconditions
  validateStateFileExist

  runContainer stop
}

installJob() {
  printInstallHeader
  verifyHostPreconditions
  createAndValidateHostStoragePath
  validateStateFileNotExist

  runContainer init

  printHeader deploying
  runContainer deploy
  createStateFile

  startChoice
}

COMMAND=""

COMMAND=help
if [ $# -gt 0 ]; then
    COMMAND=$1
fi

case $COMMAND in
  install )
    installJob
    ;;
  deploy )
    deployJob
    ;;
  start )
    startJob
    ;;
  stop )
    stopJob
    ;;
  --help|help|-h )
      echo "================================================================================================"
      echo " TiDB Enterprise ${TIDB_VERSION} Installer"
      echo "================================================================================================"
      echo ""
      echo "${0##*/} [COMMAND] {OPTIONS}"
      echo ""
      echo "Available commands:"
      echo "  install                        Install TiDB Enterprise Installer on the host, and deploy the TiDB cluster"
      echo ""
      echo "  deploy                         Deploy the TiDB cluster"
      echo ""
      echo "  start                          Start the TiDB cluster"
      echo ""
      echo "  stop                           Stop the TiDB cluster"
      echo ""
      echo "================================================================================================"
      exit 0
      ;;
  *)
      echo -e "${RED}Unknown command '$COMMAND'${NC}"
      exit $UNKNOWN_COMMAND_EXIT_CODE
    ;;
esac