#!/usr/bin/env bash

export YA_USER=$USER
export MAIN_PATH=$(dirname "$0")
export PYTHON=$(which python)
export PYTHONPATH=$MAIN_PATH/lib/python3/dist-packages/:$MAIN_PATH/lib
TOOLS_DIR="$MAIN_PATH/tools"
DOC_DIR="$MAIN_PATH/doc"
MAN_DIR="$MAIN_PATH/man"
BCC_MAN_DIR="$MAIN_PATH/man/bcc"
BPF_MAN_DIR="$MAIN_PATH/man/bpf"
IPROUTE2_MAN_DIR="$MAIN_PATH/man/iproute2"
BCC_DOC_DIR="$DOC_DIR/bcc"
BPF_DOC_DIR="$DOC_DIR/bpf"
BCC_TOOLS_DIR="$TOOLS_DIR/bcc"
BPF_TOOLS_DIR="$TOOLS_DIR/bpf"
IPROUTE2_TOOLS_DIR="$TOOLS_DIR/iproute2"

. ${MAIN_PATH}/lib/yandex.sh

usage() {
  echo """
Usage: ya tool sre [commands] <tool> [args] [opts]

Available commands:

  -h | help  - print help
       list  - list of available tools
        man  - show the manual page for the specific tool
        doc  - show the doc page for the specific tool
        bcc  - run custom bcc script
        bpftrace  - run custom bpftrace script

Examples:
	# Run bcc script
	ya tool sre runqlen

	# Run bpftrace script
	ya tool sre tcp_bl_drops.bt

	# Run custom bcc script
	ya tool sre bcc ./my_script

	# Run custom bpftrace script
	ya tool sre bpftrace ./my_script

	# Show man
	ya tool sre man runqlen

	# Show doc
	ya tool sre doc runqlen

Dirs:
	MAIN_PATH=$MAIN_PATH
	TOOLS_DIR=$TOOLS_DIR
	MAN_DIR=$MAN_DIR
	DOC_DIR=$DOC_DIR
"""
}

cmd_list() {
  echo """\
# Tools
$(find $TOOLS_DIR -type f -maxdepth 1 -printf '  %f\n' | sort)

# BCC tools:
$(find $BCC_TOOLS_DIR -type f -printf '  %f\n' | sort)

# BPF tools:
$(find $BPF_TOOLS_DIR -type f -printf '  %f\n' | sort)

# iproute2 tools:
$(find $IPROUTE2_TOOLS_DIR -type f -printf ' %f\n' | sort)
"""
}

cmd_man() {
	tool_man_path="$MAN_DIR/$1.gz"
	bcc_man_path="$BCC_MAN_DIR/$1.gz"
	bpf_man_path="$BPF_MAN_DIR/$1.gz"
	iproute2_man_path="$IPROUTE2_MAN_DIR/$1.gz"
  if [ -f $bcc_man_path ]; then
		man $bcc_man_path && exit
	elif [ -f $bpf_man_path ]; then
		man $bpf_man_path && exit
	elif [ -f $tool_man_path ]; then
		man $tool_man_path && exit
	elif [ -f $iproute2_man_path ];  then
		man $iproute2_man_path && exit
	else
		fn=${1%%.*}
		find $MAN_DIR -name $fn*.gz -exec man {} \;
	fi
}

cmd_doc() {
	tool_doc_path="$DOC_DIR/$1.txt"
	bcc_doc_path="$BCC_DOC_DIR/$1.txt"
	bpf_doc_path="$BPF_DOC_DIR/$1.txt"
  if [ -f $bcc_doc_path ]; then
		less $bcc_doc_path && exit
	elif [ -f $bpf_doc_path ]; then
		less $bpf_doc_path && exit
	elif [ -f $tool_doc_path ]; then
		less $tool_doc_path && exit
	else
		fn=${1%%.*}
		find $DOC_DIR -name $fn* -exec less {} \;
	fi
}

cmd_bcc() {
    cmd_run_sudo PYTHONPATH=$PYTHONPATH LD_LIBRARY_PATH=$MAIN_PATH/lib/ $PYTHON "$@"
}

cmd_bpf() {
		cmd_run_sudo $TOOLS_DIR/bpftrace "$@"
}

run() {
	[ -z $1 ] && usage && exit

	tool=$1
	shift
	if [ -f $BCC_TOOLS_DIR/$tool ]; then
		cmd_bcc $BCC_TOOLS_DIR/$tool "$@"
	elif [ -f $BPF_TOOLS_DIR/$tool ]; then
		cmd_bpf $BPF_TOOLS_DIR/$tool "$@"
	elif [ -f $IPROUTE2_TOOLS_DIR/$tool ]; then
		cmd_run_sudo $IPROUTE2_TOOLS_DIR/$tool "$@"
	elif [ -f $TOOLS_DIR/$tool ]; then
		cmd_run $TOOLS_DIR/$tool "$@"
	else
		ya tool $tool "$@"
	fi
}

case "$1" in
  -h|--help|help) usage;;
  list) cmd_list;;
  man) shift; cmd_man $1;;
  doc) shift; cmd_doc $1;;
  bcc) shift; cmd_bcc "$@";;
  *) run "$@";;
esac
