#
# Displays Git repository information.
#
# Authors:
#   Sorin Ionescu <sorin.ionescu@gmail.com>
#
# Usage:
#   Define the following styles in a prompt theme setup function.
#
#     # %s - Special action name (am, merge, rebase).
#     zstyle ':omz:module:git' action 'action:%s'
#
#     # %a - Indicator to notify of added files.
#     zstyle ':omz:module:git' added 'added:%a'
#
#     # %A - Indicator to notify of ahead branch.
#     zstyle ':omz:module:git' ahead 'ahead:%A'
#
#     # %B - Indicator to notify of behind branch.
#     zstyle ':omz:module:git' behind 'behind:%B'
#
#     # %b - Branch name.
#     zstyle ':omz:module:git' branch 'branch:%b'
#
#     # %C - Indicator to notify of a clean working directory.
#     zstyle ':omz:module:git' clean 'clean'
#
#     # %c - SHA-1 hash.
#     zstyle ':omz:module:git' commit 'commit:%c'
#
#     # %d - Indicator to notify of deleted files.
#     zstyle ':omz:module:git' deleted 'deleted:%d'
#
#     # %D - Indicator to notify of dirty files.
#     zstyle ':omz:module:git' dirty 'dirty:%D'
#
#     # %m - Indicator to notify of modified files.
#     zstyle ':omz:module:git' modified 'modified:%m'
#
#     # %p - HEAD position in relation to the nearest branch, remote, tag.
#     zstyle ':omz:module:git' position 'position:%p'
#
#     # %R - Remote name.
#     zstyle ':omz:moduleit' remote 'remote:%R'
#
#     # %r - Indicator to notify of renamed files.
#     zstyle ':omz:module:git' renamed 'renamed:%r'
#
#     # %S - Indicator to notify of stashed files.
#     zstyle ':omz:module:git' stashed 'stashed:%S'
#
#     # %U - Indicator to notify of unmerged files.
#     zstyle ':omz:module:git' unmerged 'unmerged:%U'
#
#     # %u - Indicator to notify of untracked files.
#     zstyle ':omz:module:git' untracked 'untracked:%u'
#
#     # Ignore submodule when it is 'dirty', 'untracked', 'all', or 'none'.
#     zstyle ':omz:module:git:ignore' submodule ''
#
#     # Prompts.
#     zstyle ':omz:module:git' info \
#       'prompt'  ' git:(%b%C%D)' \
#       'rprompt' ''
#

# Load dependencies.
omodload 'trap'

# Gets the Git special action (am, merge, rebase, etc.).
# Borrowed from vcs_info and edited.
function _git-action {
  local action=''
  local action_dir
  local git_dir="$(git-dir)"

  for action_dir in \
    "${git_dir}/rebase-apply" \
    "${git_dir}/rebase" \
    "${git_dir}/../.dotest"; do
    if [[ -d "$action_dir" ]] ; then
      if [[ -f "${action_dir}/rebasing" ]] ; then
        action='rebase'
      elif [[ -f "${action_dir}/applying" ]] ; then
        action='am'
      else
        action='am/rebase'
      fi
      print "$action"
      return 0
    fi
  done

  for action_dir in \
    "${git_dir}/rebase-merge/interactive" \
    "${git_dir}/.dotest-merge/interactive"; do
    if [[ -f "$action_dir" ]]; then
      print 'rebase-i'
      return 0
    fi
  done

  for action_dir in \
    "${git_dir}/rebase-merge" \
    "${git_dir}/.dotest-merge"; do
    if [[ -d "$action_dir" ]]; then
      print 'rebase-m'
      return 0
    fi
  done

  if [[ -f "${git_dir}/MERGE_HEAD" ]]; then
    print 'merge'
    return 0
  fi

  if [[ -f "${git_dir}/CHERRY_PICK_HEAD" ]]; then
    print 'cherry-pick'
    return 0
  fi

  if [[ -f "${git_dir}/BISECT_LOG" ]]; then
    print 'bisect'
    return 0
  fi

  return 1
}

# Turns off git-info for the current repository.
function _git-info-abort {
  if ! is-true "$_git_info_executing"; then
    return 1
  fi

  cat >&2 <<EOF


Gathering status for certain repositories is time intensive.
By pressing CTRL + C, you have turned off prompt Git status
for this repository.

To revert, execute:
  git-info on

EOF

  unset _git_info_executing
  git config --bool prompt.showinfo false
  git-info
  return 0
}
add-zsh-trap INT _git-info-abort

# Gets the Git status information.
function git-info {
  # Extended globbing is needed to parse repository status.
  setopt LOCAL_OPTIONS
  setopt EXTENDED_GLOB

  local action
  local action_format
  local action_formatted
  local added=0
  local added_format
  local added_formatted
  local ahead
  local ahead_and_behind
  local ahead_and_behind_cmd
  local ahead_format
  local ahead_formatted
  local ahead_or_behind
  local behind
  local behind_format
  local behind_formatted
  local branch
  local branch_format
  local branch_formatted
  local branch_info
  local clean
  local clean_formatted
  local commit
  local commit_format
  local commit_formatted
  local deleted=0
  local deleted_format
  local deleted_formatted
  local dirty=0
  local dirty_format
  local dirty_formatted
  local ignore_submodule
  local -A info_formats
  local info_format
  local line_number=0
  local modified=0
  local modified_format
  local modified_formatted
  local position
  local position_format
  local position_formatted
  local remote
  local remote_cmd
  local remote_format
  local remote_formatted
  local renamed=0
  local renamed_format
  local renamed_formatted
  local stashed=0
  local stashed_format
  local stashed_formatted
  local status_cmd
  local unmerged=0
  local unmerged_format
  local unmerged_formatted
  local untracked=0
  local untracked_format
  local untracked_formatted

  # Clean up previous $git_info.
  unset git_info
  typeset -gA git_info

  # Return if not inside a Git repository work tree.
  if ! is-true "$(git rev-parse --is-inside-work-tree 2> /dev/null)"; then
    return 1
  fi

  if (( $# > 0 )); then
    if [[ "$1" == [Oo][Nn] ]]; then
      git config --bool prompt.showinfo true
    elif [[ "$1" == [Oo][Ff][Ff] ]]; then
      git config --bool prompt.showinfo false
    else
      print "usage: $0 [ on | off ]" >&2
    fi
    return 0
  fi

  # Return if git-info is disabled.
  if ! is-true "${$(git config --bool prompt.showinfo):-true}"; then
    return 1
  fi

  # Used to abort and turn git-info off on SIGINT.
  _git_info_executing=true

  # Use porcelain status for easy parsing.
  status_cmd='git status --porcelain'

  # Gets the remote name.
  remote_cmd='git rev-parse --symbolic-full-name --verify HEAD@{upstream}'

  # Gets the commit difference counts between local and remote.
  ahead_and_behind_cmd='git rev-list --count --left-right HEAD...@{upstream}'

  # Ignore submodule status.
  zstyle -s ':omz:module:git:ignore' submodule 'ignore_submodule'
  if [[ -n "$ignore_submodule" ]]; then
    status_cmd+=" --ignore-submodules=${ignore_submodule}"
  fi

  # Format commit.
  commit="$(git rev-parse HEAD 2> /dev/null)"
  if [[ -n "$commit" ]]; then
    zstyle -s ':omz:module:git' commit 'commit_format'
    zformat -f commit_formatted "$commit_format" "c:$commit"
  fi

  # Format stashed.
  if [[ -f "$(git-dir)/refs/stash" ]]; then
    stashed="$(git stash list 2> /dev/null | wc -l | awk '{print $1}')"
    zstyle -s ':omz:module:git' stashed 'stashed_format'
    zformat -f stashed_formatted "$stashed_format" "S:$stashed"
  fi

  # Format action.
  action="$(_git-action)"
  if [[ -n "$action" ]]; then
    zstyle -s ':omz:module:git' action 'action_format'
    zformat -f action_formatted "$action_format" "s:$action"
  fi

  # Get current status.
  while IFS=$'\n' read line; do
    # Count added, deleted, modified, renamed, unmerged, untracked, dirty.
    # T (type change) is undocumented, see http://git.io/FnpMGw.
    # For a table of scenarii, see http://i.imgur.com/2YLu1.png.
    [[ "$line" == ([ACDMT][\ MT]|[ACMT]D)\ * ]] && (( added++ ))
    [[ "$line" == [\ ACMRT]D\ * ]] && (( deleted++ ))
    [[ "$line" == ?[MT]\ * ]] && (( modified++ ))
    [[ "$line" == R?\ * ]] && (( renamed++ ))
    [[ "$line" == (AA|DD|U?|?U)\ * ]] && (( unmerged++ ))
    [[ "$line" == \?\?\ * ]] && (( untracked++ ))
    (( dirty++ ))
  done < <(${(z)status_cmd} 2> /dev/null)

  # Format branch.
  branch="${$(git symbolic-ref -q HEAD)##refs/heads/}"
  if [[ -n "$branch" ]]; then
    zstyle -s ':omz:module:git' branch 'branch_format'
    zformat -f branch_formatted "$branch_format" "b:$branch"

    # Format remote.
    remote="${$(${(z)remote_cmd} 2> /dev/null)##refs/remotes/}"
    if [[ -n "$remote" ]]; then
      zstyle -s ':omz:module:git' remote 'remote_format'
      zformat -f remote_formatted "$remote_format" "R:$remote"

      # Get ahead and behind counts.
      ahead_and_behind="$(${(z)ahead_and_behind_cmd} 2> /dev/null)"

      # Format ahead.
      ahead="$ahead_and_behind[(w)1]"
      if (( $ahead > 0 )); then
        zstyle -s ':omz:module:git' ahead 'ahead_format'
        zformat -f ahead_formatted "$ahead_format" "A:$ahead"
      fi

      # Format behind.
      behind="$ahead_and_behind[(w)2]"
      if (( $behind > 0 )); then
        zstyle -s ':omz:module:git' behind 'behind_format'
        zformat -f behind_formatted "$behind_format" "B:$behind"
      fi
    fi
  else
    # Format position.
    position="$(git describe --contains --all HEAD 2> /dev/null)"
    if [[ -n "$position" ]]; then
      zstyle -s ':omz:module:git' position 'position_format'
      zformat -f position_formatted "$position_format" "p:$position"
    fi
  fi

  # Format added.
  if (( $added > 0 )); then
    zstyle -s ':omz:module:git' added 'added_format'
    zformat -f added_formatted "$added_format" "a:$added_format"
  fi

  # Format deleted.
  if (( $deleted > 0 )); then
    zstyle -s ':omz:module:git' deleted 'deleted_format'
    zformat -f deleted_formatted "$deleted_format" "d:$deleted_format"
  fi

  # Format modified.
  if (( $modified > 0 )); then
    zstyle -s ':omz:module:git' modified 'modified_format'
    zformat -f modified_formatted "$modified_format" "m:$modified"
  fi

  # Format renamed.
  if (( $renamed > 0 )); then
    zstyle -s ':omz:module:git' renamed 'renamed_format'
    zformat -f renamed_formatted "$renamed_format" "r:$renamed"
  fi

  # Format unmerged.
  if (( $unmerged > 0 )); then
    zstyle -s ':omz:module:git' unmerged 'unmerged_format'
    zformat -f unmerged_formatted "$unmerged_format" "U:$unmerged"
  fi

  # Format untracked.
  if (( $untracked > 0 )); then
    zstyle -s ':omz:module:git' untracked 'untracked_format'
    zformat -f untracked_formatted "$untracked_format" "u:$untracked"
  fi

  # Format dirty and clean.
  if (( $dirty > 0 )); then
    zstyle -s ':omz:module:git' dirty 'dirty_format'
    zformat -f dirty_formatted "$dirty_format" "D:$dirty"
  else
    zstyle -s ':omz:module:git' clean 'clean_formatted'
  fi

  # Format info.
  zstyle -a ':omz:module:git' info 'info_formats'
  for info_format in ${(k)info_formats}; do
    zformat -f REPLY "$info_formats[$info_format]" \
      "A:$ahead_formatted" \
      "B:$behind_formatted" \
      "D:$dirty_formatted" \
      "R:$remote_formatted" \
      "S:$stashed_formatted" \
      "U:$unmerged_formatted" \
      "a:$added_formatted" \
      "b:$branch_formatted" \
      "C:$clean_formatted" \
      "c:$commit_formatted" \
      "d:$deleted_formatted" \
      "m:$modified_formatted" \
      "p:$position_formatted" \
      "r:$renamed_formatted" \
      "s:$action_formatted" \
      "u:$untracked_formatted"
    git_info[$info_format]="$REPLY"
  done

  unset REPLY
  unset _git_info_executing

  return 0
}

git-info "$@"