# Copyright (c) 2024 Kana Steimle
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

_shell="$(basename -- "$0")"
if [ "$(echo "$_shell" | cut -c1)" = "-" ]; then
	_shell="$(echo "$_shell" | cut -c2-)"
fi
if [ "$COLUMNS" = "" ]; then
	COLUMNS="$(tput cols 2>&1)"
	if [ "$COLUMNS" = "" ]; then
		COLUMNS=48
	fi
fi

pcol_pwd='1;34'
pcol_command='1;34'
pcol_error='1;31'
pcol_time='1;30'

pcol_hostname='0;35'
pcol_username='0;36'
pcol_prompt=$pcol_hostname
prompt_label='$'

if [ "$(id -u)" -eq 0 ]; then
	pcol_pwd='1;33'
	pcol_username='1;31'
	pcol_prompt='0;31'
	pcol_command='1;31'
	prompt_label='#'
else
	for group in $(id -Gn); do
		if [ "$group" = "wheel" ]; then
			pcol_username='1;32'
			pcol_command='1;36'
			#prompt_label='$'
			break
		fi
	done
fi

# _pc: Print color
if [ "$_shell" = "oksh" ] || [ "$_shell" = "rksh" ]; then
	_pc() {
		printf '%b' "\\[\033[$1m\]"
	}
elif [ "$_shell" = "rbash" ] || [ "$_shell" = "bash" ] || [ "$_shell" = "dash" ] || [ "$_shell" = "mksh" ] || [ "$_shell" = "sh" ]; then
	_pc() {
		printf '%b' '\033['"$1m"
	}
else
	_pc() {
		:
	}
fi

# _pt: Print title
_pt() {
	printf '%b' '\033]0;'"$@\007" >&2
}

pcol_shell="$pcol_hostname"
if [ "$_shell" = "rksh" ] || [ "$_shell" = "rbash" ]; then
	pcol_shell='0;31'
elif [ "$_shell" = "mksh" ]; then
	pcol_shell='0;32'
elif [ "$_shell" = "bash" ]; then
	pcol_shell='0;33'
elif [ "$_shell" = "dash" ]; then
	pcol_shell='0;34'
elif [ "$_shell" = "oksh" ]; then
	pcol_shell='1;35'
elif [ "$_shell" = "sh" ]; then
	pcol_shell='1;37'
fi

_prompterror() {
	e=$?
	[ $e != 0 ] && printf "%s" "$(_pc $pcol_error)$e | "
}
_prompttime() {
	_pc $pcol_time
	printf "%s" "$(_pc $pcol_time)[$(date +%H:%M:%S)]"
}
_promptuser() {
	_pc $pcol_username
	printf "%s" "$(id -un)"
}
_prompthost() {
	_pc $pcol_hostname
	printf "%s" "@$(hostname)"
}
_promptpath() {
	_pc $pcol_pwd
	local printpwd="${PWD#"$HOME"}"
	[ ${#PWD} -ne ${#printpwd} ] && printpwd="~$printpwd"
	local max="$(($COLUMNS / 2 - 32))"
	if [ "$max" -lt 24 ]; then max=24; fi
	while [ "${#printpwd}" -gt $max ]; do
		printpwd=".../$(echo "$printpwd" | cut -d/ -f3-)"
	done
	echo "$printpwd"
	_pt "$printpwd"
}
_promptshell() {
	_pc $pcol_shell
	printf "%s" "$_shell "
}
_promptprompt() {
	_pc $pcol_prompt
	printf "%s" "$prompt_label"
}
_promptcmd() {
	_pc $pcol_command
}

# Automatically construct a prompt based on a bitmask
setprompt() {
	bitmask="$1"
	if [ "$bitmask" = "" ]; then
		# Predict desired features of prompt based on screen width
		if [ "$COLUMNS" -lt 32 ]; then
			bitmask=0
		elif [ "$COLUMNS" -lt 48 ]; then
			bitmask=1
		elif [ "$COLUMNS" -lt 64 ]; then
			bitmask=3
		elif [ "$COLUMNS" -lt 80 ]; then
			bitmask=7
		else
			bitmask=15
		fi
	fi
	PS1='$(_prompterror)'
	if [ $(((bitmask) / 8 % 2)) -eq 1 ]; then
		PS1="$PS1"'$(_prompttime) '
	fi
	if [ $(((bitmask) / 4 % 2)) -eq 1 ]; then
		PS1="$PS1"'$(_promptuser)$(_prompthost) '
	fi
	if [ $(((bitmask) / 2 % 2)) -eq 1 ]; then
		PS1="$PS1"'$(_promptpath) '
	fi
	if [ $(((bitmask) / 1 % 2)) -eq 1 ]; then
		PS1="$PS1"'$(_promptshell)'
	fi
	PS1="$PS1"'$(_promptprompt) $(_promptcmd)'
}

setprompt