#!/usr/bin/env bash
# Copyright 2018 PingCAP, 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,
# See the License for the specific language governing permissions and
# limitations under the License.

PROG=${0##*/}

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

DB_HOST="127.0.0.1"
DB_PORT="4000"
DB_USER="*****"
DB_PSWD="*****"
DB_SCHE="mysql"
SILENT=0
TOP_N=0

TMP_FILE="/tmp/${PROG}.tmp"

function usage() {
  cat 1>&2 <<-EOF
USAGE: $PROG [OPTIONS] [DB_NAME]

DB_NAME Database to use, equivalent to the -D flag if given.

OPTIONS:
  -h Connect to host of tidb-server.
  -P Port number to use for connection, built-in default (4000).
  -u User for login.
  -p Password to use when connecting to tidb-server.
  -D Database to use.
  -s Be more silent, don't print header and footer.
  -n list the top N of table size.

EOF
}

function err_exit() {
  echo "ERROR: $1" 1>&2
  exit 1
}

function pre_checks() {
  which mysql >/dev/null 2>&1 || err_exit "mysql client not installed"
  which curl >/dev/null 2>&1 || err_exit "curl not found"
  which printf >/dev/null 2>&1 || err_exit "printf not found"
}

function mysql_query() {
  MYSQL_PWD="$DB_PSWD" mysql -h $DB_HOST -P$DB_PORT -u$DB_USER -D $DB_SCHE -Ns -e "$1"
}

function table_space_usage() {
  for tbl in $(mysql_query "show tables"); do
    usage=$(curl -s http://$DB_HOST:10080/tables/$DB_SCHE/$tbl/disk-usage)
    re='^[0-9]+$'
    if [[ $usage =~ $re ]]; then
      echo $tbl $usage
    else
      echo $tbl 0
    fi
  done
}

pre_checks

# main from here

while getopts ":h:P:u:D:p:sn:" opt; do
  case "$opt" in
    h  ) DB_HOST=$OPTARG;;
    P  ) DB_PORT=$OPTARG;;
    u  ) DB_USER=$OPTARG;;
    D  ) DB_SCHE=$OPTARG;;
    p  ) DB_PSWD="$OPTARG";;
    s  ) SILENT=1;;
    n  ) TOP_N="$OPTARG";;
    \? ) usage; err_exit "Unknown option: -$OPTARG";;
    :  )
      if [[ $OPTARG == "p" ]]; then
        echo -e "Enter password: \c"
        read PSWD
        DB_PSWD="-p$PSWD"
      else
        usage; err_exit "Missing option arugment for -$OPTARG"
      fi
      ;;
    *  ) usage; exit 1;;
  esac
done

shift $(($OPTIND - 1))

if [[ $# -ge 1 ]]; then
  DB_SCHE=$1
fi

DB_VERSION=$(mysql_query "select version()")
[[ -z $DB_VERSION ]] && err_exit "connect to tidb-server $DB_HOST $DB_PORT $DB_USER $DB_SCHE failed"
table_space_usage > $TMP_FILE
total=$(cat $TMP_FILE | awk '{t+=$2}END{print t}')
if [[ $TOP_N -gt 0 ]]; then
  cat $TMP_FILE | sort -k2 -n -r | head -$TOP_N > $TMP_FILE.2
else
  cat $TMP_FILE | sort -k2 -n -r > $TMP_FILE.2
fi

if [[ $SILENT -eq 0 ]]; then
  printf "## TiDB Version:  %s\n" $DB_VERSION
  printf "## Database Name: %s\n" $DB_SCHE
  printf "## Total Size:    %s MB\n" $total

  printf "\n%-48s %14s\n" "TABLE" "SPACE USAGE"
  echo "-----------------------------------------------------------------"
fi

cat $TMP_FILE.2 | while read -r line; do
printf "%-48s %11d MB\n" $line
done

if [[ $SILENT -eq 0 ]]; then
  real_total=$(cat $TMP_FILE.2 | awk '{t+=$2}END{print t}')
  echo "-----------------------------------------------------------------"
  printf "%-48s %11d MB\n" "TOTAL SIZE" $real_total
fi

# cleanup
rm -rf $TMP_FILE*