# -*- sh -*-

zmodload zsh/pcre

autoload colors && colors

_VCS_RESET="%{${reset_color}%}"
_VCS_COLOUR="%{${reset_color}${fg_bold[cyan]}%}"
_VCS_BRANCH_COLOUR="%{${reset_color}${fg_bold[yellow]}%}"
_SUDO_RESET="%{${reset_color}%}"
_SUDO_ALLOWED="%{${reset_color}${fg_bold[red]}%}"

_VCS_GENERIC_AWK_SCRIPT=$(cat <<'EOF'
BEGIN {
  unk=0;
  add=0;
  del=0;
  mod=0;
  con=0;
  mer=0;
  ren=0;
  prop=0;
  other=0;
  others="";
}

{
  if ($1 == "?") { unk++; } else
  if ($1 == "A") { add++; } else
  if ($1 == "M") { mod++; } else
  if ($1 == "D" || $1 == "-D") { del++; } else
  if ($1 == "C") { con++; } else
  if ($1 == "N" || $1 == "+N") { add++; } else
  if ($1 == "P" || $1 == "P.") { mer++; } else
  if ($1 == "R") { ren++; } else
  if ($1 == "Performing") { } else
  if ($1 == "RM") { ren++; mod++; } else
  if ($1 == "*") { prop++; } else
  if ($1 == "X") { } else
  if ($1 == "") { } else
  if ($1 == "Current") { } else
  { other++; others = others $1}
}

END {
  ret="";
  if (unk > 0) ret = ret "?";
  if (add > 0) ret = ret "A";
  if (del > 0) ret = ret "D";
  if (ren > 0) ret = ret "R";
  if (mod > 0) ret = ret "M";
  if (prop > 0) ret = ret "F"
  if (con > 0) ret = ret "C[" con "]";
  if (mer > 0) ret = ret "+P";
  if (other > 0) ret = ret "-";
  print ret others;
}
EOF
)

_VCS_GIT_AWK_SCRIPT=$(cat <<'GITEOF'
BEGIN {
  unk=0;
  add=0;
  del=0;
  mod=0;
  con=0;
  mer=0;
  ren=0;
  other=0;
  others="";
  mode=0;
}

/^\?\?/ {
  unk++;
}

/^A/ {
  add++;
}

/^.A/ {
  add++;
}

/^R/ {
  ren++;
}

/^.R/ {
  ren++;
}

/^D/ {
  del++;
}

/^.D/ {
  del++;
}

/^M/ {
  mod++;
}

/^.M/ {
  mod++;
}

END {
  ret="";
  if (unk > 0) ret = ret "?";
  if (add > 0) ret = ret "A";
  if (del > 0) ret = ret "D";
  if (ren > 0) ret = ret "R";
  if (mod > 0) ret = ret "M";
  if (con > 0) ret = ret "C[" con "]";
  if (mer > 0) ret = ret "+P";
  if (other > 0) ret = ret "-"; 
  print ret others;
}
GITEOF
)

_vcs_status () {
    # Parse a command (passed in as $@) 's output like an svn st or bzr st -S
    # Decide on the working tree status of the tree if possible
    local vcs_status
    _VCS_STATUS=""
    vcs_status=$("$@" "${_VCS_ROOT}" 2>/dev/null)
    # If there are no changes, return now
    [[ "x$vcs_status" == "x" ]] && return 0
    # Okay there are changes
    _VCS_STATUS=$(echo "${vcs_status}" | awk "${_VCS_GENERIC_AWK_SCRIPT}")
}

_vcs_git_status () {
    # Parse a command (passed in as $@) 's output like git status
    # Decide on the working tree status of the tree if possible
    local vcs_status
    vcs_status=$("$@" 2>/dev/null)
    _VCS_STATUS=$(echo "${vcs_status}" | awk "${_VCS_GIT_AWK_SCRIPT}")
}

_find_bzr_root () {
    # Scanning from CWD upward, look for a .bzr directory
    # once found, echo that top level
    local cdir
    cdir=$PWD
    while [[ $(dirname $cdir) != $cdir ]]; do
	[[ -d $cdir/.bzr && -r $cdir/.bzr/README ]] && {
	    echo $cdir
	    return 0
	}
	cdir=$(dirname $cdir)
    done
    return 1
}

# Find if we're in a BZR tree and if so, store its data and succeed
# otherwise return failure
_find_bzr () {
    local bzr_root
    bzr_root=$(_find_bzr_root) || return 1
    # Okay we know it's a bzr tree, let's gather data and form a
    # bit of prompt data...
    
    if ! test -d "$bzr_root/.bzr/branch"; then
        _VCS="bzr-repo"
        _VCS_ROOT="$bzr_root"
        _VCS_BRANCH="N/A"
        _VCS_REVNO="N/A"
        _VCS_STATUS=""
	return 0
    fi

    # Check if we're bound, have push/pull/merge locations saved etc.

    _VCS="bzr"
    _VCS_ROOT="$bzr_root"
    _VCS_BRANCH=$(bzr nick)
    local tree_rev branch_rev
    tree_rev=$(bzr revno --tree || echo "No Tree")
    branch_rev=$(bzr revno)
    if test "x$tree_rev" = "x$branch_rev"; then
        _VCS_REVNO="$branch_rev"
    else
        _VCS_REVNO="$tree_rev/$branch_rev"
    fi
    _vcs_status bzr status -S
}

_svn_uri () {
    # echo the uri of the provided svn dir
    test -r $1/.svn/entries || return 0
    svn info --non-interactive $1 | grep "^URL:" | cut -d' ' -f2
}

_svn_revno () {
    svn info --non-interactive $1 | grep "^Revision:" | cut -d' ' -f2
}

# Find if we're in an SVN tree and if so, find the uppermost level which
# is a full prefix of the CWD
_find_svn () {
    local cwd_uri parent_uri this_uri this_dir
    # Do nothing if we're not in an SVN tree
    test -r .svn/entries || return 1
    # Okay, we're in an SVN tree, so we're going to look upwards until we
    # find a dir which either isn't svn or isn't a strict parent of where
    # we are now.
    cwd_uri=$(_svn_uri $(pwd))
    parent_uri=$(_svn_uri $(pwd)/..)
    this_uri=${cwd_uri}
    this_dir=$(pwd)
    latest_ver=$(_svn_revno $this_dir)
    earliest_ver=$(_svn_revno $this_dir)
    while [[ "x$parent_uri" = "x$(dirname $this_uri)" ]]; do
	this_dir=$(dirname $this_dir)
	this_uri=$parent_uri
	parent_uri=$(_svn_uri ${this_dir}/..)
	this_ver=$(_svn_revno $this_dir)
	if [ "0$this_ver" -gt "0$latest_ver" ]; then
	    latest_ver="$this_ver"
	fi
	if [ "0$this_ver" -lt "0$earliest_ver" ]; then
	    earliest_ver="$this_ver"
	fi
    done
    # Now, this_dir is the top level (the VCS root)
    _VCS="svn"
    _VCS_ROOT="$this_dir"
    if [ "0$latest_ver" != "0$earliest_ver" ]; then
	_VCS_REVNO="$earliest_ver:$latest_ver"
    else
	_VCS_REVNO="$latest_ver"
    fi
    _vcs_status svn status
    # Branch is a little harder.
    [[ "$this_uri" -pcre-match trunk || "$this_uri" -pcre-match branch ]] || {
	# No mention of banches or trunks, just use the basename of the uri
	_VCS_BRANCH=$(basename $this_uri)
	return 0
    }
    [[ "$this_uri" -pcre-match trunk/([^/]+) || "$this_uri" -pcre-match ([^/]+)/trunk ]] && {
	# trunk/blah or blah/trunk
	local trunkname
	trunkname=$(basename $this_dir)
	if [[ "x$trunkname" = "x$match" ]]; then
	    _VCS_BRANCH="trunk"
	else
	    _VCS_BRANCH="trunk($match)"
	fi
	return 0
    }
    [[ "$this_uri" -pcre-match branch[^/]*/(.+)$ ]] && {
	# branches/*
	_VCS_BRANCH="$match"
	return 0
    }

    # Otherwise, use the basename of the svn uri
    _VCS_BRANCH=$(basename $this_uri)
}

_find_git_root () {
    # Scanning from CWD upward, look for a .git directory
    # once found, echo that top level
    local cdir
    cdir=$PWD
    while [[ $(dirname $cdir) != $cdir ]]; do
	[[ -d $cdir/.git && -r $cdir/.git/config ]] && {
	    echo $cdir
	    return 0
	}
	cdir=$(dirname $cdir)
    done
    return 1
}

# Retrieve the git branch name
_get_git_branch () {
    local branchname
    branchname=$(git symbolic-ref -q HEAD 2>/dev/null ||
	git name-rev --name-only HEAD)
    echo "${branchname#refs/heads/}"
}

# Retrieve the seven digit revcode for the branch
_get_git_revcode () {
    (git show-ref -s --abbrev=7 $(_get_git_branch) || git rev-parse --short $(_get_git_branch)) | head -1
}

# Find if we're in a GIT tree, and fill out _VCS_* if we are.
_find_git () {
    local git_root
    git_root=$(_find_git_root) || return 1
    _VCS="git"
    _VCS_ROOT="$git_root"
    _VCS_BRANCH=$(_get_git_branch)
    _VCS_REVNO=$(_get_git_revcode)
    _vcs_git_status git status --porcelain
}

_vcsify_path () {
    # Take the current path and VCSify it given the contents of the _VCS* variables
    _VCS_PROMPT_ROOT=$(print -D $_VCS_ROOT)
    if [[ "x${_VCS_STATUS}" != "x" ]]; then
	_VCS_STATUS=",${_VCS_STATUS}"
    fi
    if [[ $(basename $_VCS_ROOT) == $_VCS_BRANCH ]]; then
	_VCS_PROMPT_BRANCH="[${_VCS_REVNO}${_VCS_STATUS}]"
    else
	_VCS_PROMPT_BRANCH="[${_VCS_BRANCH}:${_VCS_REVNO}${_VCS_STATUS}]"
    fi
    _VCS_PROMPT_LEAF=${PWD#${_VCS_ROOT}}
}

_find_cargo_root () {
    # Scanning from CWD upward, look for a Cargo.toml file
    # once found, echo that top level
    local cdir
    cdir=$PWD
    while [[ $(dirname $cdir) != $cdir ]]; do
	[[ -r $cdir/Cargo.toml ]] && {
	    echo $cdir
	    return 0
	}
	cdir=$(dirname $cdir)
    done
    return 1
}

_RUST_BADGE_CRAB="🦀"
_RUST_BADGE_STABLE="🏠"
_RUST_BADGE_BETA="β"
_RUST_BADGE_NIGHTLY="🌙"

_find_rust () {
    local rust_root
    local toolchain
    rust_root=$(_find_cargo_root) || return 1
    toolchain=$(cd "$rust_root"; rustup show active-toolchain)
    case "$toolchain" in
	stable-*)
	    _RUST_LABEL="${_RUST_BADGE_CRAB}${_RUST_BADGE_STABLE}"
	    ;;
	beta-*)
	    _RUST_LABEL="${_RUST_BADGE_CRAB}${_RUST_BADGE_BETA}"
	    ;;
	nightly-*)
	    _RUST_LABEL="${_RUST_BADGE_CRAB}${_RUST_BADGE_NIGHTLY}"
	    ;;
	*)
	    _RUST_LABEL="${_RUST_BADGE_CRAB}${toolchain}"
    esac

    return 0
}

_find_debian_source_dir () {
    # Scanning from CWD upward, look for a debian/changelog file
    # once found, echo that top level
    local cdir
    cdir=$PWD
    while [[ $(dirname $cdir) != $cdir ]]; do
	[[ -d $cdir/debian && -r $cdir/debian/changelog ]] && {
	    echo $cdir
	    return 0
	}
	cdir=$(dirname $cdir)
    done
    return 1
}

_DEBIAN_BADGE="🍥"

_find_debian () {
    local debian_root
    local version
    debian_root=$(_find_debian_source_dir) || return 1
    version=$(cd "${debian_root}"; dpkg-parsechangelog -S Version)
    _DEBIAN_LABEL="${_DEBIAN_BADGE}${version}"
}

_KEY_BADGE="🔑"

_check_keys () {
    if ssh-add -l >/dev/null 2>/dev/null; then
	_SSH_LABEL="${_KEY_BADGE}"
    else
	_SSH_LABEL=""
    fi
}

# Now hide any functions which won't work due to missing apps.

which bzr >/dev/null 2>&1 || { 
    _find_bzr () { return 1; }
}

which svn >/dev/null 2>&1 || {
    _find_svn () { return 1; }
}

which git-name-rev >/dev/null 2>&1 || test -x /usr/lib/git-core/git-name-rev || {
    _find_git () { return 1; }
}

which dpkg-parsechangelog >/dev/null 2>&1 || {
    _find_debian () { return 1; }
}

# Sudo support

SUDO=$(which sudo)
if test -x "$SUDO" >/dev/null 2>/dev/null; then
    could_sudo () { sudo -n -l /bin/true >/dev/null 2>/dev/null; }
else
    # No sudo support
    could_sudo () { return 1; }
fi

sudo_precmd () {
    if could_sudo; then
	_SUDO_COLOUR="yes"
    else
	_SUDO_COLOUR=""
    fi
}

# And put it all together

precmd () {
    local ok
    ok=1
    _find_svn ||
    _find_bzr ||
    _find_git ||
    {
	ok=0
    }
    psvar=()
    title_precmd
    sudo_precmd
    [[ $ok == 1 ]] && {
	_vcsify_path
	psvar[1]="${_VCS}"
	psvar[2]="${_VCS_PROMPT_ROOT}"
	psvar[3]="${_VCS_PROMPT_BRANCH}"
	psvar[4]="${_VCS_PROMPT_LEAF}"
    }
    psvar[5]="${_SUDO_COLOUR}"

    _LABELS=""
    _find_rust &&   { _LABELS="${_LABELS}${_RUST_LABEL}" }
    _find_debian && { _LABELS="${_LABELS}${_DEBIAN_LABEL}" }

    psvar[6]="$_LABELS"

    _check_keys

    psvar[7]="${_SSH_LABEL}"
}

_VCS_PROMPT="${_VCS_COLOUR}(%1v)${_VCS_RESET}"
_VCS_RPROMPT="%2v${_VCS_BRANCH_COLOUR}%3v${_VCS_RESET}%4v"

SHELLDEPTH=$((${SHELLDEPTH:-0} + 1))
export SHELLDEPTH

if test "x$SHELLDEPTH" != "x1"; then
    _SHELLDEPTH="%{${reset_color}${fg_bold[blue]}%}(${SHELLDEPTH})%{${reset_color}%}"
else
    _SHELLDEPTH=""
fi

# Prompt as machine%                                           /path/on/right
PS1="%(5V.${_SUDO_ALLOWED}.${_SUDO_RESET})%m${_SUDO_RESET}${_SHELLDEPTH}%(1V.${_VCS_PROMPT}.)%(6V.(%6v).)%# "
RPS1="%(1V.${_VCS_RPROMPT}.%~)%7v"

if [ "x$debian_chroot" = "x" -a -e /etc/debian_chroot ]; then
  debian_chroot=$(cat /etc/debian_chroot)
fi

if [ "x$debian_chroot" != "x" ]; then
  PS1="%(5V.${_SUDO_ALLOWED}.${_SUDO_RESET})%m${_SUDO_RESET}/$debian_chroot${_SHELLDEPTH}%(1V.${_VCS_PROMPT}.)%(6V.(%6v).)%# "
fi