#!/bin/bash
# shebang for syntax highlighting purposes

# Colors
_NON="\[\033[0m\]"
_BLD="\[\033[1m\]"
_BLK="\[\033[30m\]"
_RED="\[\033[31m\]"
_GRN="\[\033[32m\]"
_YLW="\[\033[33m\]"
_BLU="\[\033[34m\]"
_PRP="\[\033[35m\]"
_CYN="\[\033[36m\]"
_WHT="\[\033[37m\]"
_BACK_BLK="\[\033[40m\]"
_BACK_RED="\[\033[41m\]"
_BACK_GRN="\[\033[42m\]"
_BACK_YLW="\[\033[43m\]"
_BACK_BLU="\[\033[44m\]"
_BACK_PRP="\[\033[45m\]"
_BACK_CYN="\[\033[46m\]"
_BACK_WHT="\[\033[47m\]"
_BRIGHT_BLK="\[\033[90m\]"
_BRIGHT_RED="\[\033[91m\]"
_BRIGHT_GRN="\[\033[92m\]"
_BRIGHT_YLW="\[\033[93m\]"
_BRIGHT_BLU="\[\033[94m\]"
_BRIGHT_PRP="\[\033[95m\]"
_BRIGHT_CYN="\[\033[96m\]"
_BRIGHT_WHT="\[\033[97m\]"
_BRIGHT_BACK_BLK="\[\033[100m\]"
_BRIGHT_BACK_RED="\[\033[101m\]"
_BRIGHT_BACK_GRN="\[\033[102m\]"
_BRIGHT_BACK_YLW="\[\033[103m\]"
_BRIGHT_BACK_BLU="\[\033[104m\]"
_BRIGHT_BACK_PRP="\[\033[105m\]"
_BRIGHT_BACK_CYN="\[\033[106m\]"
_BRIGHT_BACK_WHT="\[\033[107m\]"

# Number of guix packages to show before truncating in `guix shell`
_GUIX_MAXLIST=1

# Minimum microseconds to show a count
_MIN_US=0

# Number of characters of the {start,end} of the IP you SSH into
_SSH_LIMIT_START=4
_SSH_LIMIT_END=4

# Colors (higher level)
_COLOR_SSH_IP="${_CYN}"
_COLOR_SSH_PORT="${_CYN}"
_COLOR_SSH_PORT_SEP="${_PRP}${_BLD}"
_COLOR_SSH_ASTRSK="${_WHT}"
_COLOR_SSH_BRACKS="${_BLU}"

_COLOR_GUIX_BRACKS="${_BLD}${_BLU}"
_COLOR_GUIX_GUIX="${_YLW}"
_COLOR_GUIX_SEP=""
_COLOR_GUIX_PKGS="${_WHT}"
_COLOR_GUIX_OUT="${_WHT}"
_COLOR_GUIX_OUT_SEP="${_WHT}"
_COLOR_GUIX_PLUS=""
_COLOR_GUIX_PKGNUM="${_PRP}"
_COLOR_GUIX_PKG=""

_COLOR_UHD_BRACKS="${_BLD}${_BLU}"
_COLOR_UHD_ROOT="${_BLD}${_RED}"
_COLOR_UHD_NONROOT="${_BLD}${_YLW}"
_COLOR_UHD_AT="${_BLD}${_PRP}"
_COLOR_UHD_HOST="${_BLD}${_CYN}"
_COLOR_UHD_SEP="${_BLD}${_PRP}"
_COLOR_UHD_DIR="${_BLD}${_GRN}"

_COLOR_CMDTIME="${_YLW}"

__debug_trap() {
        # Set necessary pre-command variables (PROMPT_COMMAND is a
        # command so its excluded here)
        if [[ "$BASH_COMMAND" != "$PROMPT_COMMAND"
                && "$LAST_BASH_COMMAND" == "$PROMPT_COMMAND" ]]; then
                INC_TIME=1
                LAST_RT="$EPOCHREALTIME" # This should be the last
                                           # thing done in this function
        fi
        LAST_BASH_COMMAND="$BASH_COMMAND" # But this needs to be
}

trap '__debug_trap' DEBUG

__get_cmd_time() {
        # Set hours minutes seconds and remove preceding zeros
        local YEAR
        local DAYS
        local HOUR
        local MINS
        local SECS
        ((YEAR=10#$(($(TZ=UTC printf '%(%Y)T' "$CMD_TIME")))-1970))
        ((DAYS=10#$(TZ=UTC printf '%(%j)T' "$CMD_TIME")))
        ((HOUR=10#$(TZ=UTC printf '%(%H)T' "$CMD_TIME")))
        ((MINS=10#$(TZ=UTC printf '%(%M)T' "$CMD_TIME")))
        ((SECS=10#$(TZ=UTC printf '%(%S)T' "$CMD_TIME")))

        # Choose whether or not to print hours minutes and seconds
        [[ $CMD_TIME -ge 31536000 ]] && printf '%sy ' "${YEAR}"
        [[ $CMD_TIME -ge 86400    ]] && printf '%sd ' "${DAYS}"
        [[ $CMD_TIME -ge 3600     ]] && printf '%sh ' "${HOUR}"
        [[ $CMD_TIME -ge 60       ]] && printf '%sm ' "${MINS}"
        [[ $CMD_TIME -ge 1        ]] && printf '%ss ' "${SECS}"
        [[ $CMD_US   -ge $_MIN_US ]] && printf '%sμs' "${CMD_US:-0}"
}

__sig() {
        # Giant switch case for getting the name of the signal (`kill -l`)
        local ITR=0
        local RET
        for RET in "$@"; do
                if [[ $ITR != 0 ]]; then
                        printf '%s|%s' "$_WHT" "$_RED"
                fi
                ((ITR++))
                case $RET in
                        126): "ACCES"  ;;
                        127): "NOENT"  ;;
                        129): "HUP"    ;;
                        130): "INT"    ;;
                        131): "QUIT"   ;;
                        132): "ILL"    ;;
                        133): "TRAP"   ;;
                        134): "ABRT"   ;;
                        135): "BUS"    ;;
                        136): "FPE"    ;;
                        137): "KILL"   ;;
                        138): "USR"1   ;;
                        139): "SEGV"   ;;
                        140): "USR"2   ;;
                        141): "PIPE"   ;;
                        142): "ALRM"   ;;
                        143): "TERM"   ;;
                        144): "STKFLT" ;;
                        145): "CHLD"   ;;
                        146): "CONT"   ;;
                        147): "STOP"   ;;
                        148): "TSTP"   ;;
                        149): "TTIN"   ;;
                        150): "TTOU"   ;;
                        151): "URG"    ;;
                        152): "XCPU"   ;;
                        153): "XFSZ"   ;;
                        154): "VTALRM" ;;
                        155): "PROF"   ;;
                        156): "WINCH"  ;;
                        157): "IO"     ;;
                        158): "PWR"    ;;
                        159): "SYS"    ;;
                        16[3-9]|1[7-8][0-9]|19[0-2]): "RT$((RET-128))" ;; # Savagery
                        *):   "$RET"   ;; # Print exit code if not in list
                esac
                printf '%s' "$_"
        done
}

__ssh() {
        local CON=($SSH_CONNECTION)
        local SRV_IP="${CON[2]}"
        [[ $_SSH_LIMIT_START -lt 0 ]] && _SSH_LIMIT_START=0
        [[ $_SSH_LIMIT_END -lt 0 ]] && _SSH_LIMIT_END=0
        [[ -z "$SRV_IP" ]] && return
        local SRV_PORT="${CON[3]}"
        local SRV_IP_CUT="${_NON}${_COLOR_SSH_IP}${SRV_IP}"
        [[ ${#SRV_IP} -gt $((_SSH_LIMIT_START+_SSH_LIMIT_END)) ]] && local SRV_IP_CUT="${_NON}${_COLOR_SSH_IP}${SRV_IP:0:$_SSH_LIMIT_START}${_COLOR_SSH_ASTRSK}*${_NON}${_COLOR_SSH_IP}${SRV_IP: -$_SSH_LIMIT_END}"

        printf '%s' "${_NON}${_COLOR_SSH_BRACKS}[${_NON}${SRV_IP_CUT}${_COLOR_SSH_PORT_SEP}:${_NON}${_COLOR_SSH_PORT}${SRV_PORT}${_COLOR_SSH_BRACKS}]${_NON}"
}

__guix_env() {
        local re='[\(]?\("([0-9a-z+_:-]*)"$'
        local re1='"([a-z]*)"$'
        if [[ "$GUIX_ENVIRONMENT" ]]; then
                local line
                local PACKAGES=($(while read -r line; do
                                [[ "$line" =~ $re ]] && printf '%s ' "${BASH_REMATCH[1]}"
                        done < $GUIX_ENVIRONMENT/manifest))
                local i=0
                local OUTPUTS=($(while read -r line; do
                                if [[ "$line" =~ $re1 ]]; then
                                        local match="${BASH_REMATCH[1]}"
                                        if [[ ! "${PACKAGES[$i]}" == "$match" ]]; then
                                                printf '%s ' "$match"
                                        fi
                                        ((i++))
                                fi
                        done < $GUIX_ENVIRONMENT/manifest))
                local FULL_PKGS=($(for i in $(seq 0 $((${#PACKAGES[@]}-1))); do
                                        if [[ "${OUTPUTS[$i]}" == "out" ]]; then
                                                printf '%s ' "${_COLOR_GUIX_PKGS}${PACKAGES[$i]}${_NON}"
                                        else
                                                printf '%s ' "${_COLOR_GUIX_PKGS}${PACKAGES[$i]}${_NON}${_COLOR_GUIX_OUT_SEP}:${_NON}${_COLOR_GUIX_OUT}${OUTPUTS[$i]}${_NON}"
                                        fi
                                done))
                local FULL_PKG_STR="${FULL_PKGS[@]:0:$_GUIX_MAXLIST}"
                local EXTRA="$((${#FULL_PKGS[@]}-_GUIX_MAXLIST))"
                if [ ${_GUIX_MAXLIST} -gt 0 ] && [ $EXTRA -gt 0 ]; then
                        FULL_PKG_STR+=" ${_NON}${_COLOR_GUIX_PLUS}+ "
                fi
                if [ $EXTRA -gt 1 ]; then
                        FULL_PKG_STR+="${_NON}${_COLOR_GUIX_PKGNUM}$EXTRA ${_NON}${_COLOR_GUIX_PKG}pkgs"
                elif [ $EXTRA -eq 1 ]; then
                        FULL_PKG_STR+="${_NON}${_COLOR_GUIX_PKGNUM}1 ${_NON}${_COLOR_GUIX_PKG}pkg"
                fi

                printf %s "${_NON}${_COLOR_GUIX_BRACKS}[${_NON}${_COLOR_GUIX_GUIX}Guix${_NON}${_COLOR_GUIX_SEP}: ${_NON}${FULL_PKG_STR}${_COLOR_GUIX_BRACKS}] "
        fi
}

__prompt() {
        # Get exit code (must be first)
        local PLC=("${PIPESTATUS[@]}")

        # Reset time when prompt was first displayed after command
        # this contributes to the ~40 microsecond difference in $CMD_US and the actual time it took
        if [[ "$INC_TIME" != 0 ]]; then 
                PROMPT_RT="$EPOCHREALTIME"
                INC_TIME=0
        fi

        # *_RT may not be set
        LAST_RT="${LAST_RT:-$EPOCHREALTIME}"
        PROMPT_RT="${PROMPT_RT:-$EPOCHREALTIME}"

        # Get relative times

        # Remove decimal point, simulating multiplying by 1 million
        PROMPT_RT1M="${PROMPT_RT/.}"
        PROMPT_RT1M="${PROMPT_RT1M/,}"
        LAST_RT1M="${LAST_RT/.}"
        LAST_RT1M="${LAST_RT1M/,}"

        ((CMD_US=PROMPT_RT1M-LAST_RT1M))

        CMD_TIME=0

        # Divide by 1 million to get a more accurate difference
        ((CMD_TIME=CMD_US/1000000))

        [[ ${#CMD_US} -lt 6 ]] || CMD_US="${CMD_US: -6}"
        ((CMD_US=10#$CMD_US))
        [[ $CMD_US -lt 0 ]] && CMD_US=0

        # Set prompt sections

        local SIG
        local COL
        local UHD
        local TIME
        local IND
        local GUIX

        # Text

        # ssh detection and indicator
        [[ "$SSH_CONNECTION" ]] && local SSH="$(__ssh) "

        # [INT], [4], etc.
        local RET
        for RET in "${PLC[@]}"; do
                if [[ $RET -gt 0 ]]; then
                        SIG="$(printf '%s[%s%s%s] ' "$_BLU" "$_RED" "$(__sig "${PLC[@]}")" "$_BLU")"
                        break
                fi
        done

        # Guix shell environment
        GUIX="$(__guix_env)"

        # [user@homeserver:~]
        COL="$([[ $UID == 0 ]] && printf '%s' "${_COLOR_UHD_ROOT}" || printf '%s' "${_COLOR_UHD_NONROOT}")"
        UHD="${_NON}${_COLOR_UHD_BRACKS}[${_NON}${COL}\u${_NON}${_COLOR_UHD_AT}@${_NON}${_COLOR_UHD_HOST}\h${_NON}${_COLOR_UHD_SEP}:${_NON}${_COLOR_UHD_DIR}\w${_NON}${_COLOR_UHD_BRACKS}]"

        # 2y 351d 12m 43s 382969μs
        TIME="${_NON}${_COLOR_CMDTIME}$(__get_cmd_time)"

        # Random colored $ or #
        _RAND256="\[\033[38;2;$((RANDOM%256));$((RANDOM%256));$((RANDOM%256))m\]"
        IND="${_RAND256}\\$"

        # Set the prompt
        PS1="${_NON}${SSH}${GUIX}${SIG}${UHD} ${TIME} ${IND} ${_NON}"
}
        
PROMPT_COMMAND=__prompt