autoload -Uz promptinit
setopt appendhistory
setopt extendedglob
setopt correctall
bindkey -v
bindkey -v '^?' backward-delete-char

## --------------------------------------------------
## Helper functions
## --------------------------------------------------

exists () {
    command -v "$1" &> /dev/null
}

alias_if_N () {
    while [ "$1" != "--" ]; do
         exists "$1" || return 1
         shift
    done
    shift
    for a in "$@"; do
	alias "$a"
    done
}

alias_if () {
    alias_if_N "$1" -- "${@:2}"
}

alias_if_else () {
    if ! alias_if_N "$1" -- "$2"; then
        alias "$3"
    fi
}

alias_if_sudo () {
    alias_if_N "$1" -- "${@:2}"
}

## --------------------------------------------------
# autocompletion
## --------------------------------------------------

# load stuff (order matters)
zmodload zsh/complist
autoload -Uz compinit && compinit -i
autoload -Uz bashcompinit && bashcompinit

# use vi-like menu for selection
bindkey -M menuselect 'h' vi-backward-char
bindkey -M menuselect 'k' vi-up-line-or-history
bindkey -M menuselect 'j' vi-down-line-or-history
bindkey -M menuselect 'l' vi-forward-char
bindkey -M menuselect 'gg' beginning-of-history
bindkey -M menuselect 'G' end-of-history
bindkey -M menuselect 'gj' vi-forward-blank-word
bindkey -M menuselect 'gk' vi-backward-blank-word
bindkey -M menuselect '^[[Z' reverse-menu-complete
zstyle ':completion:*' menu select

# add nice colors for descriptions/warnings/messages
zstyle ':completion:*:*:*:*:descriptions' format '%F{green}-- %d --%f'
zstyle ':completion:*:messages' format ' %F{purple} -- %d --%f'
zstyle ':completion:*:warnings' format ' %F{red}-- no matches found --%f'

# use sane grouping
zstyle ':completion:*' group-name ''

# make file lists useful by showing relevant info
zstyle ':completion:*' file-list all

# Highlight the current autocomplete option
zstyle ':completion:*:default' list-colors ${(s.:.)LS_COLORS}

# ignore full hostnames in ssh config file
# the default ssh function does not seem to differentiate b/t
# HOST and HOSTNAME directives
zstyle ':completion:*:(ssh|scp|rsync):*:hosts' ignored-patterns '*(.|:)*' loopback ip6-loopback localhost ip6-localhost broadcasthost

## --------------------------------------------------
# syntax highlighting a la fish
## --------------------------------------------------

hlpath="$XDG_DATA_HOME/zsh/plugins/zsh-syntax-highlighting/zsh-syntax-highlighting.zsh"
if [[ -e "$hlpath" ]]; then
    . "$hlpath"
fi

## --------------------------------------------------
# history
## --------------------------------------------------

setopt hist_expire_dups_first
setopt hist_ignore_dups
setopt hist_reduce_blanks
setopt inc_append_history
setopt share_history

autoload -Uz up-line-or-beginning-search
autoload -Uz down-line-or-beginning-search
zle -N up-line-or-beginning-search
zle -N down-line-or-beginning-search
bindkey '\eOA' up-line-or-beginning-search
bindkey '\e[A' up-line-or-beginning-search
bindkey '\eOB' down-line-or-beginning-search
bindkey '\e[B' down-line-or-beginning-search

## --------------------------------------------------
## PROMPT 
## --------------------------------------------------

setopt prompt_subst

if [[ ${EUID} -eq 0 ]]; then
    user_color="%F{red}"
else
    user_color="%F{cyan}"
fi

HOSTNAME=""
if [ -n "$SSH_CLIENT" ] || [ -n "$SSH_TTY" ]; then
    HOSTNAME="@%M"
fi	

## git prompt

zsh-git-off () {
    export ZSH_USE_GIT=0
}

zsh-git-on () {
    export ZSH_USE_GIT=1
}

zsh-git-on

PROMPT=""
RPROMPT=""

git_info() {
    # Exit if not inside a Git repository
    ! git rev-parse --is-inside-work-tree > /dev/null 2>&1 && return

    # Git branch/tag, or name-rev if on detached head
    local GIT_LOCATION=$(git symbolic-ref -q HEAD || \
                             git name-rev --name-only \
                                 --no-undefined --always HEAD)
    GIT_LOCATION=${GIT_LOCATION#(refs/heads/|tags/)}

    local AHEAD="%F{red}⇡NUM%f"
    local BEHIND="%F{cyan}⇣NUM%f"
    local MERGING="%F{magenta⚡︎%f"
    local UNTRACKED="%F{red}●%f"
    local MODIFIED="%F{yellow}●%f"
    local STAGED="%F{green}●%f"

    local -a DIVERGENCES
    local -a FLAGS

    local NUM_AHEAD="$(git log --oneline @{u}.. 2> /dev/null | wc -l | tr -d ' ')"
    if [ "$NUM_AHEAD" -gt 0 ]; then
        DIVERGENCES+=( "${AHEAD//NUM/$NUM_AHEAD}" )
    fi

    local NUM_BEHIND="$(git log --oneline ..@{u} 2> /dev/null | wc -l | tr -d ' ')"
    if [ "$NUM_BEHIND" -gt 0 ]; then
        DIVERGENCES+=( "${BEHIND//NUM/$NUM_BEHIND}" )
    fi

    local GIT_DIR="$(git rev-parse --git-dir 2> /dev/null)"
    if [ -n $GIT_DIR ] && test -r $GIT_DIR/MERGE_HEAD; then
        FLAGS+=( "$MERGING" )
    fi

    if [[ -n $(git ls-files --other --exclude-standard 2> /dev/null) ]]; then
        FLAGS+=( "$UNTRACKED" )
    fi

    if ! git diff --quiet 2> /dev/null; then
        FLAGS+=( "$MODIFIED" )
    fi

    if ! git diff --cached --quiet 2> /dev/null; then
        FLAGS+=( "$STAGED" )
    fi

    local -a GIT_INFO
    [[ ${#DIVERGENCES[@]} -ne 0 ]] && GIT_INFO+=( "${(j::)DIVERGENCES}" )
    [[ ${#FLAGS[@]} -ne 0 ]] && GIT_INFO+=( "${(j::)FLAGS}" )
    GIT_INFO+=( "%F{green}$GIT_LOCATION%f" )
    echo "─(%B${(j: :)GIT_INFO}%b)"
}

function zle-line-init zle-keymap-select {
    if [[ "$ZSH_USE_GIT" -eq 1 ]]; then
       GIT=$(git_info)
    else;
       GIT=""
    fi
    if [ -z "$CONDA_DEFAULT_ENV" ]; then
       CONDA=""
    else;
       CONDA="─(%B%F{cyan}$CONDA_DEFAULT_ENV%f%b)"
    fi
    NORMAL="%F{yellow}N%f"
    INSERT="%F{cyan}I%f"
    VIMODE="─%B(${${KEYMAP/vicmd/$NORMAL}/(main|viins)/$INSERT})%b"
    USER="$user_color%n$HOSTNAME%f"
    TIME="%F{green}%*%f"
    WD="%F{green}%1~%f"
    PROMPT=$'\n%B┌($USER)─($TIME)$VIMODE$GIT$CONDA\n└─($WD)─>%b '
    zle reset-prompt
}

zle -N zle-line-init
zle -N zle-keymap-select

export KEYTIMEOUT=1

## --------------------------------------------------
## Title
## --------------------------------------------------

precmd () { echo -ne "\033]0;urxvt: ${PWD}\007" }

## --------------------------------------------------
## Aliases 
## --------------------------------------------------

setopt autopushd

## ASSUME these will always exist

alias ls='LC_COLLATE=C ls --color --group-directories-first'

alias d='dirs -v | head -10'
alias d1='cd -'
alias d2='cd -2'
alias d3='cd -3'
alias d4='cd -4'
alias d5='cd -5'
alias d6='cd -6'
alias d7='cd -7'
alias d8='cd -8'
alias d9='cd -9'

alias u='cd ..'
alias uu='cd ../..'
alias uuu='cd ../../..'
alias uuuu='cd ../../../..'
alias uuuuu='cd ../../../../..'

## diff is weird; only use color if ver >= 3.4

diffver="$(diff -v | grep -o '\([0-9]\+\.[0-9]\+\)')"
if [ ${diffver//\.*/} -ge 3 ] && [ ${diffver//*\./} -ge 4 ]; then
    alias diff='diff --color'
    alias_if zdiff zdiff='zdiff --color'
fi

## don't assume the rest of these will exist

alias_if grep \
    grep='grep --color=auto' \
    als='alias | grep' \
    hs='history 1 | grep' \
    envs='env | grep'

alias_if sudo sudo='sudo '


if ! alias_if eza \
        ll='eza -alhg --group-directories-first' \
        llt='eza -T' \
        llg='ll --git'; then
    alias ll='ls -alh'
    alias_if tree llt='tree'
fi

alias_if_sudo reboot reboot='sudo reboot'
alias_if_sudo poweroff poweroff='sudo poweroff'
alias_if_sudo halt halt='sudo halt'
alias_if_sudo shutdown shutdown='sudo shutdown'

alias_if_sudo systemctl \
    sc="sudo systemctl" \
    sce="sudo systemctl enable" \
    scd="sudo systemctl disable" \
    scs="sudo systemctl start" \
    sct="sudo systemctl stop" \
    scr="sudo systemctl restart"

alias_if systemctl \
    sca="systemctl status" \
    scu="systemctl --user" \
    scue="systemctl --user enable" \
    scud="systemctl --user disable" \
    scus="systemctl --user start" \
    scut="systemctl --user stop" \
    scur="systemctl --user restart"

alias_if git \
    g='git' \
    gcl='git clone' \
    ga='git add' \
    grm='git rm' \
    ga='git add' \
    gus='git reset HEAD' \
    gm="git merge" \
    gmv='git mv' \
    gs='git status -s -b' \
    gp='git push' \
    gpo='git push origin' \
    gpom='git push origin master' \
    gr='git remote' \
    gd='git diff' \
    gc='git commit -v' \
    gca='git commit -v -a' \
    gcm='git commit -v -m' \
    gcam="git commit -v -am" \
    gci='git commit --interactive' \
    gb='git branch' \
    gbl='git branch --list' \
    gco='git checkout' \
    glol="git log --graph --decorate --pretty=oneline --abbrev-commit" \
    glola="git log --graph --decorate --pretty=oneline --abbrev-commit --all" \
    gg="git log --graph --pretty=format:'%C(yellow)%h\\ %ad%Cred%d\\ %Creset%s%C(cyan)\\ [%cn]' --abbrev-commit --date=relative" \
    ggs="gg --stat" \
    gw="git whatchanged" \
    gt="git tag" \
    gnew="git log HEAD@{1}..HEAD@{0}" \
    ggui="git gui"

alias_if tmux \
    tls='tmux ls' \
    tks='tmux kill-session -t' \
    tas='tmux attach-session -t'

alias_if chezmoi \
    czp="chezmoi apply" \
    cze="chezmoi edit" \
    czea="chezmoi edit --apply" \
    czd="chezmoi diff" \
    cza="chezmoi add" \
    czu="chezmoi update" \
    czc="chezmoi cd"

alias_if xclip \
    xclip-img="xclip -selection clipboard -t image/png -i"

# squeue -> proxy for "slurm"
alias_if squeue \
    squ="squeue -u $(id -u)"

## --------------------------------------------------
## Manly Colors
## --------------------------------------------------

man() {
    env LESS_TERMCAP_mb=$'\E[01;31m' \
    LESS_TERMCAP_md=$'\E[01;38;5;13m' \
    LESS_TERMCAP_me=$'\E[0m' \
    LESS_TERMCAP_se=$'\E[0m' \
    LESS_TERMCAP_so=$'\E[38;5;3m' \
    LESS_TERMCAP_ue=$'\E[0m' \
    LESS_TERMCAP_us=$'\E[04;38;5;10m' \
    man "$@"
}

{{- if .desktop }}

## --------------------------------------------------
# enable x11 clipboard sync
## --------------------------------------------------
if exists xclip; then
    function x11-clip-wrap-widgets() {
        # NB: Assume we are the first wrapper and that we only wrap native widgets
        # See zsh-autosuggestions.zsh for a more generic and more robust wrapper
        local copy_or_paste=$1
        shift
    
        for widget in $@; do
            if [[ $copy_or_paste == "copy" ]]; then
                eval "
                function _x11-clip-wrapped-$widget() {
                    zle .$widget
                    xclip -in -selection clipboard <<<\$CUTBUFFER
                }
                "
            else
                eval "
                function _x11-clip-wrapped-$widget() {
                    CUTBUFFER=\$(xclip -out -selection clipboard)
                    zle .$widget
                }
                "
            fi
    
            zle -N $widget _x11-clip-wrapped-$widget
        done
    }
    
    local copy_widgets=(
        vi-yank vi-yank-eol vi-delete vi-backward-kill-word vi-change-whole-line
    )
    local paste_widgets=(
        vi-put-{before,after}
    )
    x11-clip-wrap-widgets copy $copy_widgets
    x11-clip-wrap-widgets paste  $paste_widgets
fi

{{- end }}

## --------------------------------------------------
## Python/Ruby Virtual Environments
## --------------------------------------------------

# both of these languages have similar setups for their virtual envs which
# involve a) adding a "shims" directory to PATH and then b) adding some shell
# magic to automatically call the shims

{{ if .development.python -}}
zsh-pyenv () {
    export PATH=$PYENV_ROOT/shims:$PATH

    if command -v pyenv 1>/dev/null 2>&1; then
        echo "Activating pyenv"
        eval "$(pyenv init - | sed '/PATH/d' -)"
        eval "$(pyenv virtualenv-init - | sed '/PATH/d' -)"
    else
        echo "pyenv not installed"
    fi
}
{{- end }}

{{ if .development.ruby -}}
zsh-rbenv () {
    gempaths="$(/usr/bin/gem env gempath):"
    export PATH="${gempaths//:/\/bin:}:$PATH"
    export PATH=$RBENV_ROOT/shims:$PATH

    if command -v rbenv 1>/dev/null 2>&1; then
        echo "Activating rbenv"
        eval "$(rbenv init - | sed '/PATH/d' -)"
    else
        echo "rbenv not installed"
    fi
}
{{- end }}

{{ if .development.conda -}}
zsh-conda () {
    if [ -n "$1" ]; then
        if [ -f "$1" ]; then
            source "$1"
        else
            echo "Could not source file at $1"
            return 1
        fi
    elif [ -x $XDG_DATA_HOME/mambaforge/bin/conda ]; then
        eval "$($XDG_DATA_HOME/mambaforge/bin/conda shell.zsh hook)"
    else
        echo "conda installation not found"
        return 1
    fi
    alias cdd="conda deactivate"
    alias cda="conda activate"
    alias cde="conda env"
    alias cdl="conda list"
}
{{- end }}