#!/bin/bash

# Charles Danesi, 2022.
# parts taken from https://gist.github.com/loudambiance/a41b42a4295bce6e7304

# TODO git merge status
# TODO git isLocalOnly to check if there are remotes or if this is only a local repo)
# TODO number of pushes pending
# TODO display number of modified files (maybe number of untracked as well)
# TODO display current tag
# TODO display that files are staged
# TODO display a 'ready to commit' status
# TODO display file deleted
# TODO status to display detached HEAD
# TODO status icon to indicate stash

#=========================================================
#Terminal Color Codes
#=========================================================
WHITE='\[\033[1;37m\]'
LIGHTGRAY='\[\033[0;37m\]'
GRAY='\[\033[1;49;90m\]'
BLACK='\[\033[0;30m\]'
RED='\[\033[0;31m\]'
LIGHTRED='\[\033[1;31m\]'
GREEN='\[\033[0;32m\]'
LIGHTGREEN='\[\033[1;32m\]'
ORANGE='\[\033[0;33m\]'
YELLOW='\[\033[1;33m\]'
BLUE="\[\033[0;34m\]"
LIGHTBLUE='\[\033[1;34m\]'
PURPLE='\[\033[0;35m\]'
PINK='\[\033[1;35m\]'
CYAN='\[\033[0;36m\]'
LIGHTCYAN='\[\033[1;36m\]'
DEFAULT='\[\033[0m\]'
CDNET1="\[\e[38;5;27m\]"
CDNET2="\[\e[38;5;32m\]"
DARKGRAY="\[\e[38;5;237m\]"
DARKGRAY2="\[\e[90m\]"

#=========================================================
# User Configuration
#=========================================================
# TODO refactor colors and config options.
# Configuration
#todoDir = "$HOME/code/todo.txt/"  # This MUST be set for todo.txt to work!

# Colors
clrLines=$GRAY # Lines and Arrow
clrBrackets=$GRAY # Brackets around each data item
clrError=$RED # Error displays when previous command did not return 0
clrClock=$DARKGRAY # The current time
clrScreenLabel=$CYAN # Color for terminal multiplexer label
cMPX1=$YELLOW # Color for terminal multiplexer threshold 1
cMPX2=$RED # Color for terminal multiplexer threshold 2
cBGJ=$PURPLE # Color for background job label
cBGJ1=$YELLOW # Color for background job threshold 1
cBGJ2=$RED # Color for background job threshold 2
cSTJ=$RED # Color for stopped job label
cSTJ1=$YELLOW # Color for background job threshold 1
cSTJ2=$RED # Color for  background job threshold 2
clrTaskLabel=$PURPLE # Color for todo.txt tasks label
clrTasks=$YELLOW # Color for the number of todo.txt tasks
clrSSH=$PINK # Color for brackets if session is an SSH session
clrUsername=$CDNET2 # Color of user
clrHostname=$CDNET2 # Color of hostname
clrWarnAsRoot=$RED # Color of root warning
clrWorkingDirectory=$CDNET1 # Color of current directory
clrCommand=$DEFAULT # Color of the command you type
clrSeparator=$WHITE # Color for separator for screen/jobs
clrUserSeperator=$GRAY # Color of the user and hostname separator, probably '@'
clrDefaultLines=$clrLines # default color
clrDefaultBrackets=$clrBrackets # default color
clrGitBranch=$LIGHTGREEN
clrGitHash=$CDNET1
clrCmdRunTime=$CYAN
GITBG1="\[\e[47m\]"
GITBG2="\[\e[100m\]"

# Features
fBlankLine=1  # Have a newline between previous command output and new prompt
fGitRepo=1 # Show git status
fGitStatusBar=1 # Show the git status bar above the main prompt (3 lines) NOTE: currently unused
fGitAfterCWD=1 # Show git status on 2-line prompt NOTE: currently unused
fReturnErrors=1 # Previous command return status tracker
fWindowTitle=1 # show uptime/load in titlebar
fClock=1
fShowHostname=1
fScreenTracker=1 # Terminal multiplexer tracker enabled
fSSHSession=0 # Track if session is SSH
fBackgroundJobTracker=1 # Track background jobs
fStoppedJobsTracker=1 # Track stopped jobs
fShowUserHost=1 # Show user and host  #### NOTE: CURRENTLY UNUSED ####
fWorkingDirectory=1 # Show current directory
fPathShortening=1 # Enable directory shortening (configurable)
fLastCmdRunTime=1 # Display last cmd wall time

fShowTaskCount=1 # Show task count from todo.txt, if installed

# Settings
# TODO: add thresholds for cmd wall time
# TODO: add settings for customizing todo.txt (project/context/priority display only)
shortDirectoryThreshold="2" # How many folders deep before we start trimming the path
mpxThresholdMid="0" # Terminal multiplexer threshold 1 value
mpxThresholdHigh="2" # Terminal multiplexer threshold 2 value
bgJobThresholdMid="0" # Background job threshold 1 value
bgJobThresholdHigh="2" # Background job threshold 2 value
stoppedJobThresholdMid="0" # Stopped job threshold 1 value
stoppedJobThresholdHigh="2" # Stopped job threshold 2 value

# Symbols/Indicators
userHostSeparator="@"
symScreen="s"
gitBranch="" # 
gitUntrackedFiles="?"
gitFilesAdded="A"
gitFilesRemoved="D"
gitFilesModified="M"
gitStagedChanges=""
gitUnstagedChanges="Δ"
gitStashedChanges="*"  # 
gitMerging="‼" #   
gitUncommittedChanges="✘" # staged, but not committed
gitClean="✓"
gitDirty="✘"
gitPushReady="" # 
gitBehindRemote=""
gitLocalOnly=""
gitLocalRemote=""
taskLabel="☑"

# if user-specific config exists, override default values
if [ -f $HOME/.cdprc ]; then
    . $HOME/.cdprc
fi

function promptcmd() {
    prevReturnCode=$?
    timer_stop

    # check if we're in an SSH session
    if [[ $SSH_CLIENT ]] || [[ $SSH2_CLIENT ]]; then
        lSSH_FLAG=1
    else
        lSSH_FLAG=0
    fi

    # insert newline to clear space from previous cmd
    if [ $fBlankLine -eq 1 ]; then
        PS1="\n"
    else
        PS1=
    fi

    # previous command error
    if [ $fReturnErrors -eq 1 ]; then
        if [ $prevReturnCode -ne 0 ]; then
            clrLines=$clrError
            clrBrackets=$clrError
        else
            clrLines=$clrDefaultLines
            clrBrackets=$clrDefaultBrackets
        fi
    fi

    # set window title
    if [ $fWindowTitle -eq 1 ]; then
        if [ "$STY" != "" ]; then
            PS1="\[\e]0;(`echo $STY|cut -d . -f 2`)`uptime`\a\]${PS1}"
        else
            PS1="\[\e]0;`uptime -p`, `uptime | cut -d " " -f 12-`\a\]${PS1}"
        fi
    fi

    # beginning of first line (if in a git repo)
#    if [ $fGitRepo -eq 1 ] && [ "$(is_git_repo)" == "true" ]; then
#        PS1="${PS1}  ${WHITE} ${clrGitBranch}$(git_branch)$(git_hash) $(git_untracked)\n"
#    fi

    # first line if not in a git repo
    PS1="${PS1}${debian_chroot:+($debian_chroot)}${clrLines}\342\224\214\342\224\200"

    # last cmd runtime
    if [ $fLastCmdRunTime -eq 1 ]; then
        PS1="${PS1}${clrBrackets}[${clrCmdRunTime}$timer_show${clrBrackets}]${clrLines}\342\224\200"
    fi

    # current time
    if [ $fClock -eq 1 ] ; then
        PS1="${PS1}${clrBrackets}[${clrClock}\t${clrBrackets}]${clrLines}\342\224\200"
    fi

    # detached screen sessions
    if [ $fScreenTracker -eq 1 ] ; then
        hTMUX=0
        hSCREEN=0
        mpxCounter=0
        hash tmux --help 2>/dev/null || hTMUX=1
        hash screen --version 2>/dev/null || hSCREEN=1
        if [ $hTMUX -eq 1 ] && [ $hSCREEN -eq 1 ]; then
            mpxCounter=$(echo "$(screen -ls|grep -c -i detach) + $(tmux ls 2> /dev/null | grep -c -i -v attach)" | bc)
        elif [ $hTMUX -eq 0 ] && [ $hSCREEN -eq 1 ] ; then
            mpxCounter=$(tmux ls 2>/dev/null | grep -c -i -v attach)
        elif [ $hTMUX -eq 1 ] && [ $hSCREEN -eq 0 ] ; then
            mpxCounter=$(screen -ls | grep -c -i detach)
        fi
        if [[ $mpxCounter -gt $mpxThresholdHigh ]] ; then
            PS1="${PS1}${clrBrackets}[${clrScreenLabel}${symScreen}${clrSeparator}:${cMPX2}${mpxCounter}${clrBrackets}]${clrLines}\342\224\200"
        elif [[ $mpxCounter -gt $mpxThresholdMid ]] ; then
            PS1="${PS1}${clrBrackets}[${clrScreenLabel}${symScreen}${clrSeparator}:${cMPX1}${mpxCounter}${clrBrackets}]${clrLines}\342\224\200"
        fi
    fi

    # background tasks
    if [ $fBackgroundJobTracker -eq 1 ] ; then
        BGJC=$(jobs -r|wc -l)
        if [ $BGJC -gt $bgJobThresholdHigh ] ; then
            PS1="${PS1}${clrBrackets}[${cBGJ2}&${clrSeparator}:${cBGJ2}${BGJC}${clrBrackets}]${clrLines}\342\224\200"
        elif [ $BGJC -gt $bgJobThresholdMid ] ; then
            PS1="${PS1}${clrBrackets}[${cBGJ1}&${clrSeparator}:${cBGJ2}${BGJC}${clrBrackets}]${clrLines}\342\224\200"
        fi
    fi

    # stopped jobs
    if [ $fStoppedJobsTracker -eq 1 ] ; then
        STJC=$(jobs -s | wc -l )
        if [ $STJC -gt $stoppedJobThresholdHigh ] ; then
            PS1="${PS1}${clrBrackets}[${cSTJ2}\342\234\227${clrSeparator}:${cSTJ2}${STJC}${clrBrackets}]${clrLines}\342\224\200"
        elif [ $STJC -gt $stoppedJobThresholdMid ] ; then
            PS1="${PS1}${clrBrackets}[${cSTJ1}\342\234\227${clrSeparator}:${cSTJ1}${STJC}${clrBrackets}]${clrLines}\342\224\200"
        fi
    fi

    # todo.txt task count
    if [ $fShowTaskCount -eq 1 ] ; then
        PS1="${PS1}${clrTaskLabel}"
    fi

    # user@host
    # set bracket color if in ssh session
    if [ $fSSHSession -eq 1 ] && [ $lSSH_FLAG -eq 1 ]; then
        clrSes="$clrSSH"
    else
        clrSes="$clrBrackets"
    fi

    # don't display user if root
    if [ $EUID -eq 0 ]; then
        PS1="${PS1}${clrSes}[${clrWarnAsRoot}!"
    else
        PS1="${PS1}${clrSes}[${clrUsername}\u"
    fi

    # Host Section
    if [ $fShowHostname -eq 1 ] || [ $lSSH_FLAG -eq 1 ] ; then   # Host is optional only without SSH
        PS1="${PS1}${clrUserSeperator}${userHostSeparator}${clrHostname}\h${clrSes}]${clrLines}\342\224\200"
    else
        PS1="${PS1}${clrSes}]${clrLines}\342\224\200"
    fi
    #PS1="${PS1}${userHostSeparator}${LIGHTGREEN}\h${clrSes}]${GRAY}\342\224\200"

    # current directory
    if [ $fWorkingDirectory -eq 1 ]; then
        PS1="${PS1}[${clrWorkingDirectory}\w"
        if [ $fGitRepo -eq 1 ] && [ "$(is_git_repo)" == "true" ]; then
            PS1="${PS1}${clrBrackets}]${clrLines}\342\224\200${clrBrackets}[${WHITE}${gitBranch} ${clrGitBranch}$(git_branch)$(git_hash) $(git_untracked)${clrBrackets}]"
        else
            PS1="${PS1}${clrBrackets}]"
        fi
    fi

    # prompt
    PS1="${PS1}\n${clrLines}\342\224\224\342\224\200\342\224\200> ${clrCommand}"
}

function timer_now {
    date +%s%3N
}

function timer_start {
    timer=${timer:-$(timer_now)}
}

function timer_stop {
    local d_ms=$(($(timer_now) - $timer))
    local d_s=$((d_ms / 1000))
    local ms=$((d_ms % 1000))
    local s=$((d_s % 60))
    local m=$(((d_s / 60) % 60))
    local h=$((d_s / 3600))
    if ((h > 0)); then timer_show=${h}h${m}m
    elif ((m > 0)); then timer_show=${m}m${s}s
    elif ((s >= 10)); then timer_show=${s}.$((ms / 100))s
    elif ((s > 0)); then timer_show=${s}.$((ms / 10))s
    else timer_show=${ms}ms
    fi
    unset timer
}

function is_git_repo {
	# TODO: can do this without saving as a variable
    local isREPO="$(git rev-parse --is-inside-work-tree 2>/dev/null)"
	 #if [ ! $(git remote 2>/dev/null) ]; then echo "no remote"; else echo "remote available";fi
    echo -e $isREPO
}

function git_hash {
    local current_hash="$(git log --oneline --max-count=1 2>/dev/null|cut -b -7)"
    if [ "$current_hash" != "" ]; then
	 	# TODO: change these to configurable variable
        echo -e "${YELLOW}@${clrGitHash}$current_hash"
    fi
}

# might be able to include untracked/unmerged/unstaged in one function
function git_untracked {
    local git_icons=""
    local git_status="$(git status 2>/dev/null)"
	 local git_shortStatus="$(git status --porcelain|cut -b 1-2 2>/dev/null)"
#	 local git_stash="$(git stash list 2>/dev/null)"
    # M modified
    # A added
    # D deleted
    # R renamed
    # C copied
    # byte 1 is staged, byte 2 unstaged
    #  ↑1 ↓10

	# has stash present
	# NOTE: IDK why this only works with the first '!'
	if ! git stash list &>/dev/null != ""; then
		git_icons="${git_icons}${LIGHTBLUE}${gitStashedChanges}"
	fi

    # untracked files
    #  if git status --porcelain|cut -b 1-2|grep "??" &>/dev/null ; then
	 # NOTE: this is probably the better way to chain these...
	 if [[ $git_shortStatus =~ "??" ]]; then
        git_icons="${git_icons}${RED}${gitUntrackedFiles}"
    fi
	 # added files
	 if [[ $git_shortStatus =~ "A" ]]; then
	     git_icons="${git_icons}${DEFAULT}${gitFilesAdded}"
	 fi
	 # deleted files
	 if [[ $git_shortStatus =~ "D" ]]; then
	     git_icons="${git_icons}${DEFAULT}${gitFilesRemoved}"
	 fi
	 
	 # modified/renamed/copied files
	 for i in $git_shortStatus; do
	 	case $i in
			R)
				;&
			C)
				;&
			M)
				#echo -n "${RED}${gitFilesModified}"
				git_icons="${git_icons}${DEFAULT}${gitFilesModified}"
				break;
				;;
		esac
	 done

    # unmerged ( ‼ )

	# FIXME: should not include this all in the same function

    # unstaged changes
    if git status --porcelain|cut -b 2|grep 'M\|A\|D\|R\|C' &>/dev/null; then
        git_icons="${git_icons}${LIGHTBLUE}${gitUnstagedChanges}"
    fi

    if [ "$git_icons" != "" ]; then
        git_icons="${git_icons} "
    fi
#    if git status --porcelain|cut -b 1|grep 'M\|A\|D\|R\|C' &>/dev/null; then
#        git_icons="${git_icons} ${YELLOW}✘"
#    if [[ !$git_status =~ "working tree clean" ]] ; then
#        git_icons="${git_icons}${RED}✘"
    if [[ $git_status =~ "Changes to be committed" ]]; then
        git_icons="${git_icons}${GREEN}✘" # NOTE: staged but not committed
    elif [[ $git_status =~ "Changes not staged" ]]; then
        git_icons="${git_icons}${LIGHTBLUE}✘"
    elif [[ $git_status =~ "branch is ahead" ]]; then
        git_icons="${git_icons}${YELLOW}✘" #NOTE: ahead of origin
    elif [[ $git_status =~ "nothing to commit" ]]; then
        git_icons="${git_icons}${LIGHTGREEN}${gitClean}" #NOTE: working
    else
        git_icons="${git_icons}${RED}${gitDirty}" # NOTE: working
    fi

#    if [ "$git_icons" != "" ]; then
#        echo -e " ${git_icons}"
#    fi
    echo -e $git_icons
}

function git_branch {
    local git_status="$(git status 2>/dev/null)"
    local on_branch="On branch ([^${IFS}]*)"
    local on_commit="HEAD detached at ([^${IFS}]*)"

    if [[ $git_status =~ $on_branch ]]; then
        local branch=${BASH_REMATCH[1]}
        echo "$branch"
    elif [[ $git_status =~ $on_commit ]]; then
        local commit=${BASH_REMATCH[1]}
        echo "($commit)"
    fi
}

function load_prompt() {
    # Get PIDs
    local parent_process=$(tr -d '\0' < /proc/$PPID/cmdline | cut -d \. -f 1)
    local my_process=$(tr -d '\0' < /proc/$$/cmdline | cut -d \. -f 1)

    if  [[ $parent_process == script* ]]; then
        PROMPT_COMMAND=""
        PS1="\t - \# - \u@\H { \w }\$ "
    elif [[ $parent_process == emacs* || $parent_process == xemacs* ]]; then
        PROMPT_COMMAND=""
        PS1="\u@\h { \w }\$ "
    else
        export DAY=$(date +%A)
        PROMPT_COMMAND="promptcmd"
    fi

    if [ $fPathShortening -eq 1 ]; then
	    export PROMPT_DIRTRIM=$shortDirectoryThreshold
	else
		export PROMPT_DIRTRIM=
	fi

    trap 'timer_start' DEBUG
    export PS1 PROMPT_COMMAND
}

load_prompt