# shellcheck shell=sh
# vim:filetype=sh:

# Copyright (c) 2025 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.

# This is a shell prompt designed to be used by a variety of shells.
# I initially wrote it for oksh, but have since expanded it to support a variety of shells.
# I have tested it with:
#  - oksh
#  - bash
#  - dash
#  - busybox ash
#  - mksh
#  - loksh
#  - ksh (the original)
# If you have another posix-compliant shell, there's a decent chance it'll work with that too.
#
# It provides the following features:
#  - Printing non-zero exit statuses of commands
#  - Printing the current user's name, color-coded according to their privleges (root, user in wheel group, user not in wheel group)
#  - Printing the current hostname
#  - Printing the time the prompt was drawn (and thus the time the last command finished)
#  - Printing a trimmed version of the current working directory (adjusted according to the width of the terminal window)
#  - Setting the title of the terminal window to the trimmed current working directory
#  - Printing a color-coded name of the shell
#  - A colored command prompt
#  - A very different default color scheme for the root user, so you can easily tell a root shell from a non-root shell.
# If you don't like the default colors I provided for you, that's okay; you can override them by setting the pcol_* variables before or after including the prompt.
#
# You can also individually disable certain components if you don't need them using the setprompt function.
# Pass in a number from 0-15, which is used as a bitmask to decide which features to enable. The highest bit corresponds to the leftmost feature, the lowest bit to the rightmost feature.
# If you don't pass an argument to setprompt, it will automatically choose what features to enable based on the width of the terminal window.

# Determine which shell we're using based on the name of the command
_shell="$(basename -- "$0")"
# Sometimes the shell name is prefixed with the - character
if [ "$(echo "$_shell" | cut -c1)" = "-" ]; then
	_shell="$(echo "$_shell" | cut -c2-)"
fi
# When you run a shell with su -c (from shadow-utils, not util-linux), you end up with $0 being su. Treat this like a posix shell.
if [ "$_shell" = "su" ]; then
	_shell=sh
fi

# Make sure we have a COLUMNS variable
if [ "$COLUMNS" = "" ]; then
	# If not set, set COLUMNS using the tput command.
	COLUMNS="$(tput cols 2>&1)"
	if [ "$COLUMNS" = "" ]; then
		# If tput didn't work, just make a reasonable guess. This number of columns shows the shell name, but nothing else.
		COLUMNS=48
	fi
fi

if [ "$(id -u)" -eq 0 ]; then
	# The root user gets a special set of default colors, so the user doesn't easily forget what they're doing.
	[ -z "$pcol_pwd" ] && pcol_pwd='1;33'
	[ -z "$pcol_username" ] && pcol_username='1;31'
	[ -z "$pcol_prompt" ] && pcol_prompt='1;91'
	[ -z "$pcol_command" ] && pcol_command='31'
	[ -z "$prompt_label" ] && prompt_label='#'
else
	# Most linux systems also have a wheel group which can execute commands as root. They get special default colors too.
	for group in $(id -Gn); do
		if [ "$group" = "wheel" ]; then
			[ -z "$pcol_username" ] && pcol_username='1;92'
			[ -z "$pcol_command" ] && pcol_command='96'
			break
		fi
	done
fi

# Set default colors that haven't already been set.
[ -z "$pcol_pwd" ] && pcol_pwd='1;94'
[ -z "$pcol_command" ] && pcol_command='97'
[ -z "$pcol_error" ] && pcol_error='1;91'
[ -z "$pcol_time" ] && pcol_time='1;90'
[ -z "$pcol_hostname" ] && pcol_hostname='95'
[ -z "$pcol_username" ] && pcol_username='1;96'
[ -z "$pcol_prompt" ] && pcol_prompt="$pcol_hostname"
[ -z "$prompt_label" ] && prompt_label='$'

# Set the shell name color according to the shell's name
if [ -z "$pcol_shell" ]; then
	case "$_shell" in
		rksh|rbash|rsh) pcol_shell='1;31' ;;
		mksh) pcol_shell='1;32' ;;
		bash) pcol_shell='1;33' ;;
		dash) pcol_shell='1;34' ;;
		loksh) pcol_shell='1;35' ;;
		oksh) pcol_shell='1;95' ;;
		ksh) pcol_shell='1;36' ;;
		ash) pcol_shell='1;90' ;;
		sh) pcol_shell='1;97' ;;
		*) pcol_shell="$pcol_hostname" ;;
	esac
fi

# _pc: Print color
if [ "$_shell" = "oksh" ] || [ "$_shell" = "rksh" ] || [ "$_shell" = "loksh" ] || [ "$_shell" = "rbash" ] || [ "$_shell" = "bash" ] || [ "$_shell" = "rsh" ] || [ "$_shell" = "ash" ] || [ -n "$BASH_VERSION" ]; then
	# oksh, bash, and busybox ash support the \[ and \] in prompts, which indicate not to count particular characters when figuring out the visible length of a line.
	# I also check BASH_VERSION, which allows these colors to work for bash when running under a different name (as sh).
	_pc() {
		printf '\\[\033[0;%bm\\]' "$1"
	}
elif [ "$_shell" = "dash" ] || [ "$_shell" = "mksh" ] || [ "$_shell" = "ksh" ] || [ -n "$OKSH_VERSION" ] || [ -n "$KSH_VERSION" ]; then
	# dash, mksh, and most ksh's don't support the \[ and \]. For dash this means nothing, but for mksh and ksh, this means it will wrap lines early.
	# If I want colors on those shells, I just have to live with that.
	# Also, apparently oksh turns off its \[ and \] support when running as sh.
	_pc() {
		printf '\033[0;%bm' "$1"
	}
else
	# If we don't know the shell, don't assume it can handle colors.
	_pc() {
		:
	}
fi

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

_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"
	printpwd="${PWD#"$HOME"}"
	[ ${#PWD} -ne ${#printpwd} ] && printpwd="~$printpwd"
	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"
	_pc 0
	printf " "
}
_promptprompt() {
	_pc "$pcol_prompt"
	printf "%s" "$prompt_label"
}
_promptcmd() {
	_pc "$pcol_command"
}

# Automatically construct a prompt based on a bitmask
# Multiple options can be passed, and they will be summed together.
# If no bitmask is passed, it will be automatically determined based on the available number of columns
setprompt() {
	if [ "$#" -gt 0 ]; then
		bitmask=0
		for n in "$@"; do
			bitmask="$((bitmask+(n)))"
		done
	else
		# 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='$(_pc 0)$(_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)'
	# Work around bash's inability to use the \[ and \] generated by _pc commands in PS1
	if [ "$_shell" = "bash" ] || [ "$_shell" = "rbash" ] || [ -n "$BASH_VERSION" ]; then
		PROMPT_COMMAND="PS1=\"$PS1\""
	fi
}

setprompt