#! /usr/bin/zsh

export __ZSH_GIT_BASEDIR=""
export __ZSH_GIT_SUBDIR=""
export __ZSH_GIT_BRANCH=""
export __ZSH_GIT_ACTION=""
export __ZSH_GIT_STATUS=""
export __ZSH_GIT_VARS_INVALID=1
export __ZSH_GIT_STATUS_INVALID=1


git_chpwd() {
  # On cd, invalidate git status in prompt
  export __ZSH_GIT_VARS_INVALID=1
}


git_preexec() {
  # On git command, invalidate git status in prompt
  case "$1" in
    git*)
      export __ZSH_GIT_VARS_INVALID=1
      ;;
  esac

  # *any* command could invalidate the repository status (new file, ...)
  export __ZSH_GIT_STATUS_INVALID=1
}


git_get_status() {
    # Return only git status
    local gitstat gitstatus

    gitstat=$(git status 2> /dev/null | grep '\(# Untracked\|# Changes\|# Changed but not updated:\)')
    gitstatus=""

    if [[ $(echo ${gitstat} | grep -c "^# Changes to be committed:$") > 0 ]]; then
      gitstatus='+'
    fi

    if [[ $(echo ${gitstat} | grep -c "^\# Changed but not updated:$") > 0 ]]; then
      gitstatus="${gitstatus}!"
    fi

    if [[ $(echo ${gitstat} | grep -c "^# Untracked files:$") > 0 ]]; then
      gitstatus="${gitstatus}?"
    fi

    echo $gitstatus
}


git_parse() {
    # psvar[3] == current action (merge, rebase, ...)
    # psvar[4] == current branch
    # psvar[5] == repository base directory
    # psvar[6] == current subdir into repository
    # psvar[7] == status (untracked, unstaged, staged)

    local git_dir ref base_dir sub_dir action branch gitstat gitstatus


    # If nothing has been invalidated
    if [[ "${__ZSH_GIT_VARS_INVALID}" == "0" && "${__ZSH_GIT_STATUS_INVALID}" == "0" ]]; then

      # reuse previous values
      psvar[3]=${__ZSH_GIT_ACTION}
      psvar[4]=${__ZSH_GIT_BRANCH}
      psvar[5]=${__ZSH_GIT_BASEDIR}
      psvar[6]=${__ZSH_GIT_SUBDIR}
      psvar[7]=${__ZSH_GIT_STATUS}

      return
    fi

    # If only status has been invalidated
    if [[ "${__ZSH_GIT_VARS_INVALID}" == "0" && "${__ZSH_GIT_STATUS_INVALID}" == "1" ]]; then

      # reuse previous values
      psvar[3]=${__ZSH_GIT_ACTION}
      psvar[4]=${__ZSH_GIT_BRANCH}
      psvar[5]=${__ZSH_GIT_BASEDIR}
      psvar[6]=${__ZSH_GIT_SUBDIR}

      export __ZSH_GIT_STATUS=$(git_get_status)
      psvar[7]=${__ZSH_GIT_STATUS}

      export __ZSH_GIT_STATUS_INVALID=0
      return
    fi


    # Git prompt variables are invalid. Update them.

    git_dir=$(git rev-parse --git-dir 2> /dev/null) || return

    base_dir=${$(readlink -f "$git_dir/..")/$HOME/'~'}

    sub_dir=$(git rev-parse --show-prefix)
    sub_dir=${sub_dir%/}
    ref=$(git symbolic-ref HEAD 2> /dev/null) || return

    action=""

    if [ -d "$git_dir/../.dotest" ]; then
      if [ -f "$git_dir/../.dotest/rebasing" ]; then
        action="-rebase"
      elif [ -f "$git_dir/../.dotest/applying" ]; then
        action="-am"
      else
        action="-am-rebase"
      fi
      branch="$ref"
    elif [ -f "$git_dir/.dotest-merge/interactive" ]; then
      action="-rebase-i"
      branch="$(cat "$git_dir/.dotest-merge/head-name")"
    elif [ -d "$git_dir/.dotest-merge" ]; then
      action="-rebase-m"
      branch="$(cat "$git_dir/.dotest-merge/head-name")"
    elif [ -f "$git_dir/MERGE_HEAD" ]; then
      action="-merge"
      branch="$ref"
    else
      test -f "$git_dir/BISECT_LOG" && psvar[3]="bisect"
      branch="$(git symbolic-ref HEAD 2>/dev/null)" || \
        branch="$(git describe --exact-match HEAD 2>/dev/null)" || \
        branch="$(cut -c1-7 "$git_dir/HEAD")..."
    fi

    # Status
    gitstatus=`git_get_status`

    # Got here, we're in git
    psvar[3]=${action}
    psvar[4]=${branch#refs/heads/}
    psvar[5]=${base_dir}
    psvar[6]=${sub_dir}
    psvar[7]=${gitstatus}
    
    # Save for next time
    export __ZSH_GIT_BASEDIR="${base_dir}"
    export __ZSH_GIT_SUBDIR="${sub_dir}"
    export __ZSH_GIT_BRANCH="${branch#refs/heads/}"
    export __ZSH_GIT_ACTION="${action}"
    export __ZSH_GIT_STATUS="${gitstatus}"

    export __ZSH_GIT_VARS_INVALID=0
    export __ZSH_GIT_STATUS_INVALID=0
}


chpwd_functions+='git_chpwd'
preexec_functions+='git_preexec'
precmd_functions+='git_parse'