# ============================================================
# Custom OpenBSD ksh prompt
# ============================================================


# ------------------------------------------------------------
# Colors
# ------------------------------------------------------------

RESET="$(printf '\033[0m')"

RED="$(printf '\033[31m')"
GREEN="$(printf '\033[32m')"
YELLOW="$(printf '\033[33m')"
BLUE="$(printf '\033[34m')"
MAGENTA="$(printf '\033[35m')"
CYAN="$(printf '\033[36m')"
WHITE="$(printf '\033[37m')"

# 256-color dark gray
DARK_GRAY="$(printf '\033[38;5;238m')"

BOLD_RED="$(printf '\033[1;31m')"
BOLD_GREEN="$(printf '\033[1;32m')"
BOLD_YELLOW="$(printf '\033[1;33m')"
BOLD_BLUE="$(printf '\033[1;34m')"
BOLD_CYAN="$(printf '\033[1;36m')"


# ------------------------------------------------------------
# User / root
# ------------------------------------------------------------

PROMPT_USER="${USER:-$(id -un)}"

if [ "$(id -u)" -eq 0 ]; then
    USER_COLOR="$BOLD_RED"
    PROMPT_CHAR="#"
else
    USER_COLOR="$BOLD_GREEN"
    PROMPT_CHAR="$"
fi


# ------------------------------------------------------------
# Hostname
# ------------------------------------------------------------

PROMPT_HOST="$(hostname)"
PROMPT_HOST="${PROMPT_HOST%%.*}"


# ------------------------------------------------------------
# Git
#
# Examples:
#
#   git:main ✓
#   git:feature/auth +2 ~3 ?1
#   git:main ↑2 ↓1
#
#   + = staged
#   ~ = modified
#   ? = untracked
#   ↑ = commits ahead of upstream
#   ↓ = commits behind upstream
# ------------------------------------------------------------

git_prompt()
{
    typeset branch
    typeset counts
    typeset staged modified untracked
    typeset upstream
    typeset ahead behind

    # Git must exist.
    whence git >/dev/null 2>&1 || return

    # We must be inside a Git repository.
    git rev-parse --is-inside-work-tree >/dev/null 2>&1 || return

    # Current branch.
    branch="$(git symbolic-ref --quiet --short HEAD 2>/dev/null)"

    # Detached HEAD.
    if [ -z "$branch" ]; then
        branch="@$(git rev-parse --short HEAD 2>/dev/null)"
    fi

    # Count staged, modified and untracked files.
    counts="$(
        git status --porcelain 2>/dev/null |
        awk '
        {
            index_status = substr($0, 1, 1)
            work_status  = substr($0, 2, 1)

            if (index_status == "?" && work_status == "?") {
                untracked++
            }
            else {
                if (index_status != " " && index_status != "?")
                    staged++

                if (work_status != " " && work_status != "?")
                    modified++
            }
        }

        END {
            printf "%d %d %d", staged, modified, untracked
        }
        '
    )"

    set -- $counts

    staged="${1:-0}"
    modified="${2:-0}"
    untracked="${3:-0}"

    ahead=0
    behind=0

    # Compare with upstream branch, if one exists.
    upstream="$(
        git rev-list \
            --left-right \
            --count \
            'HEAD...@{upstream}' \
            2>/dev/null
    )"

    if [ -n "$upstream" ]; then
        set -- $upstream

        ahead="${1:-0}"
        behind="${2:-0}"
    fi

    # Branch name.
    print -nr -- " ${MAGENTA}git:${CYAN}${branch}${RESET}"

    # Working tree state.
    if [ "$staged" -eq 0 ] &&
       [ "$modified" -eq 0 ] &&
       [ "$untracked" -eq 0 ]; then

        print -nr -- " ${GREEN}✓${RESET}"

    else

        if [ "$staged" -gt 0 ]; then
            print -nr -- " ${GREEN}+${staged}${RESET}"
        fi

        if [ "$modified" -gt 0 ]; then
            print -nr -- " ${YELLOW}~${modified}${RESET}"
        fi

        if [ "$untracked" -gt 0 ]; then
            print -nr -- " ${CYAN}?${untracked}${RESET}"
        fi
    fi

    # Ahead of upstream.
    if [ "$ahead" -gt 0 ]; then
        print -nr -- " ${GREEN}↑${ahead}${RESET}"
    fi

    # Behind upstream.
    if [ "$behind" -gt 0 ]; then
        print -nr -- " ${RED}↓${behind}${RESET}"
    fi
}


# ------------------------------------------------------------
# Python virtual environment
#
# Example:
#
#   py:.venv
# ------------------------------------------------------------

venv_prompt()
{
    typeset venv

    if [ -n "${VIRTUAL_ENV:-}" ]; then

        venv="${VIRTUAL_ENV##*/}"

        print -nr -- \
            " ${BLUE}py:${venv}${RESET}"
    fi
}


# ------------------------------------------------------------
# SSH indicator
#
# Normal SSH sessions have SSH_* environment variables.
#
# "su -" removes those variables, so who -m is used as a
# fallback to determine whether this terminal originated from
# a remote login.
# ------------------------------------------------------------

ssh_prompt()
{
    if [ -n "${SSH_CONNECTION:-}" ] ||
       [ -n "${SSH_CLIENT:-}" ] ||
       [ -n "${SSH_TTY:-}" ]; then

        print -nr -- \
            " ${DARK_GRAY}[SSH]${RESET}"

        return
    fi

    # Fallback after "su -".
    if who -m 2>/dev/null | grep -q '('; then

        print -nr -- \
            " ${DARK_GRAY}[SSH]${RESET}"
    fi
}


# ------------------------------------------------------------
# Build prompt
# ------------------------------------------------------------

build_prompt()
{
    typeset exit_code="$1"
    typeset prompt_dir

    prompt_dir="$PWD"

    # Replace $HOME with ~.
    case "$prompt_dir" in

        "$HOME")
            prompt_dir="~"
            ;;

        "$HOME"/*)
            prompt_dir="~${prompt_dir#$HOME}"
            ;;
    esac


    # --------------------------------------------------------
    # First line
    # --------------------------------------------------------

    # user@hostname
    print -nr -- \
        "${USER_COLOR}${PROMPT_USER}@${PROMPT_HOST}${RESET}"

    # Current directory
    print -nr -- \
        " ${BOLD_BLUE}${prompt_dir}${RESET}"

    # Git information
    git_prompt

    # Python virtual environment
    venv_prompt

    # SSH indicator
    ssh_prompt

    # Previous command failed.
    if [ "$exit_code" -ne 0 ]; then

        print -nr -- \
            " ${BOLD_RED}[exit:${exit_code}]${RESET}"
    fi


    # --------------------------------------------------------
    # Second line
    #
    # Root:
    #   #
    #
    # Normal user:
    #   $
    # --------------------------------------------------------

    print -nr -- "
${USER_COLOR}${PROMPT_CHAR}${RESET} "
}


# Don't let Python virtualenv modify PS1 itself.
export VIRTUAL_ENV_DISABLE_PROMPT=1


# ------------------------------------------------------------
# Enable prompt
# ------------------------------------------------------------

PS1='$(build_prompt $?)'