# vim: ft=sh

# Explain prompt
ARG=$1
if [[ $ARG == 'h' || $ARG == '-h' || $ARG == 'help' || $ARG == '--help' || $ARG == '?' ]]; then
    ARG="-h"
elif [[ $ARG != '' ]]; then
    echo "Invalid argument: $ARG"
    exit 0
fi

# If directory was just changed to the root of a git repository, print onefetch
if [ -n "$DIRCHANGED" ]; then
    if [ -d "$PWD/.git" ]; then
	if [[ "$(which onefetch)" == "" ]]; then
	    echo "This is a git directory, but you don't have 'onefetch' in your PATH"
        else
            onefetch
	fi
    fi
    unset DIRCHANGED
fi

# Default colours
# 0 default, 1 red 2 green 3 yellow 4 blue 5 magenta 6 cyan 7 white (grey)
PC="0" # prompt default
TC="3" # time
EC="7" # venv
CC="1" # chroot
UC="2" # user
RC="1" # root user
HC="4" # host
SC="1" # ssh host
DC="6" # dir
GC="5" # git

# Declare the SET name (as used by `setterm`) and PS string for colour options
# Reference them later like: ${COLOURS["5","SET"]} or ${COLOURS["2","PS"]}
declare -A COLOURS='(
    [0,SET]="default"
    [0,PS]="\[\033[0;0m\]"
    [1,SET]="red"
    [1,PS]="\[\033[0;31m\]"
    [2,SET]="green"
    [2,PS]="\[\033[0;32m\]"
    [3,SET]="yellow"
    [3,PS]="\[\033[0;33m\]"
    [4,SET]="blue"
    [4,PS]="\[\033[0;34m\]"
    [5,SET]="magenta"
    [5,PS]="\[\033[0;35m\]"
    [6,SET]="cyan"
    [6,PS]="\[\033[0;36m\]"
    [7,SET]="white"
    [7,PS]="\[\033[0;37m\]"
)'

# Get current git branch
function parse_git_branch() {
  BRANCH=$(git branch 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/\1/')
    if [ ! "${BRANCH}" == "" ]
    then
        STAT=$(parse_git_dirty)
        echo "${COLOURS[\"$GC\",\"PS\"]}(${BRANCH}${STAT})${COLOURS[\"0\",\"PS\"]}"
    else
        echo ""
    fi
}

# Get git status
function parse_git_dirty {
    status=$(git status 2>&1 | tee)
    dirty=$(echo -n "${status}" 2> /dev/null | grep "modified:" &> /dev/null; echo "$?")
    untracked=$(echo -n "${status}" 2> /dev/null | grep "Untracked files" &> /dev/null; echo "$?")
    ahead=$(echo -n "${status}" 2> /dev/null | grep "Your branch is ahead of" &> /dev/null; echo "$?")
    newfile=$(echo -n "${status}" 2> /dev/null | grep "new file:" &> /dev/null; echo "$?")
    renamed=$(echo -n "${status}" 2> /dev/null | grep "renamed:" &> /dev/null; echo "$?")
    deleted=$(echo -n "${status}" 2> /dev/null | grep "deleted:" &> /dev/null; echo "$?")
    bits=''
    if [ "${renamed}" == "0" ]; then
        bits=">${bits}"
    fi
    if [ "${ahead}" == "0" ]; then
        bits="*${bits}"
    fi
    if [ "${newfile}" == "0" ]; then
        bits="+${bits}"
    fi
    if [ "${untracked}" == "0" ]; then
        bits="?${bits}"
    fi
    if [ "${deleted}" == "0" ]; then
        bits="x${bits}"
    fi
    if [ "${dirty}" == "0" ]; then
        bits="!${bits}"
    fi
    if [ "${bits}" != "" ]; then
        echo " ${bits}"
    else
        echo ""
    fi
}

# Comment - allow for command history to be copied directly to bash script
PS1="${COLOURS["$PC","PS"]}# "
if [ -f "/run/.containerenv" ]; then
        # Sigil to indicate host, distrobox or ssh
        PS1="${PS1}${COLOURS["0","PS"]}⧉ "
elif [[ -z $SSH_CLIENT ]]; then
        PS1="${PS1}${COLOURS["1","PS"]}⌂ "
else
        PS1="${PS1}${COLOURS["1","PS"]}🖧 "
fi

# Time
PS1="${PS1}${COLOURS["$TC","PS"]}\\t "
if [[ $ARG == '-h' ]]; then
    setterm --foreground ${COLOURS["$TC","SET"]}; echo "time"; setterm --foreground default
fi

# Python venv
if [ "$VIRTUAL_ENV" ]; then
    VENV="${VIRTUAL_ENV##*/}"
    VENV=$(echo "$VENV" | sed -r 's/.*/[\0]/')
    PS1="${PS1}${COLOURS["$EC","PS"]}${VENV} "
    if [[ $ARG == '-h' ]]; then
         setterm --foreground ${COLOURS["$EC","SET"]}; echo "venv"; setterm --foreground default
    fi
fi

# Debian Chroot
if [ -z "$debian_chroot" ] && [ -r /etc/debian_chroot ]; then
    PS1="${PS1}${COLOURS["$CC","PS"]}<$(cat /etc/debian_chroot)>"
    if [[ $ARG == '-h' ]]; then
        setterm --foreground ${COLOURS["$CC","SET"]}; echo "debian chroot"; setterm --foreground default
    fi
fi

# User
if [[ $(whoami) == 'root' ]]; then
    UC=$RC
fi
PS1="${PS1}${COLOURS["$UC","PS"]}\\u"
if [[ $ARG == '-h' ]]; then
    setterm --foreground ${COLOURS["$UC","SET"]}; echo "user"; setterm --foreground default
fi

# @
PS1="${PS1}${COLOURS["$PC","PS"]}@"

# Host
if [[ $SSH_TTY != '' ]]; then
    HC="${SC}" # Host colour
    if [[ $ARG == '-h' ]]; then
        setterm --foreground ${COLOURS["$HC","SET"]}; echo "host"; setterm --foreground default
    fi
else
    if [[ $ARG == '-h' ]]; then
        setterm --foreground ${COLOURS["$HC","SET"]}; echo "host (ssh)"; setterm --foreground default
    fi
fi
PS1="${PS1}${COLOURS["$HC","PS"]}\\h"

# :
PS1="${PS1}${COLOURS["$PC","PS"]}:"

# Dir
PS1="${PS1}${COLOURS["$DC","PS"]}\\w "
if [[ $ARG == '-h' ]]; then
    setterm --foreground ${COLOURS["$DC","SET"]}; echo "directory"; setterm --foreground default
fi

# Git
GIT=$(parse_git_branch)
if [[ $GIT != '' ]]; then
    PS1="${PS1}${COLOURS["$GC","PS"]}${GIT}"
    if [[ $ARG == '-h' ]]; then
        setterm --foreground ${COLOURS["$GC","SET"]}; echo "git branch status"; setterm --foreground default
    fi
fi

# \n - start input on new line, again to support copy-paste into bash script
PS1="${PS1}${COLOURS["$PC","PS"]}\\n"

export PS1