autoload -Uz add-zsh-hook
autoload -Uz colors && colors
autoload -Uz edit-command-line && zle -N edit-command-line

# function venv_prompt() {
#     [[ -n ${VIRTUAL_ENV} ]] || return
#     echo -e "%F{yellow}%f ${VIRTUAL_ENV:t:gs/%/%%}"
# }

function reset_prompt() {
    for f in chpwd "${chpwd_functions[@]}" precmd "${(@)precmd_functions:#-z4h-direnv-hook}"; do
        [[ "${+functions[$f]}" == 0 ]] || "$f" &>/dev/null || true
    done
    zle reset-prompt
}

# Cursor Shape
function zle-keymap-select() {
    case $KEYMAP in
        vicmd) echo -ne '\e[1 q';;
        viins|main) echo -ne '\e[5 q';;
    esac
}
zle -N zle-keymap-select

function zle-line-init() {
    echo -ne "\e[5 q"
}
zle -N zle-line-init

function write_title_wd() {
    echo -ne "\e]0;$(print -P "%n@%m [%3~]")\007"
}
add-zsh-hook precmd write_title_wd

function reset_beam() {
    echo -ne '\e[5 q'
}
add-zsh-hook preexec reset_beam

function fancy-ctrl-z () {
    if [[ $#BUFFER -eq 0 ]]; then
        BUFFER="fg"
        zle accept-line -w
    else
        zle push-input -w
        zle clear-screen -w
    fi
}
zle -N fancy-ctrl-z

function sudo-command-line() {
    [[ -z $BUFFER ]] && zle up-history
    if [[ $BUFFER == sudo\ * ]]; then
        LBUFFER="${LBUFFER#sudo }"
    elif [[ $BUFFER == $EDITOR\ * ]]; then
        LBUFFER="${LBUFFER#$EDITOR }"
        LBUFFER="sudoedit $LBUFFER"
    elif [[ $BUFFER == sudoedit\ * ]]; then
        LBUFFER="${LBUFFER#sudoedit }"
        LBUFFER="$EDITOR $LBUFFER"
    else
        LBUFFER="sudo $LBUFFER"
    fi
}
zle -N sudo-command-line

function lf-wrap() {
    LF_DIRFILE="$(mktemp -u)"

    env lfX -last-dir-path="$LF_DIRFILE" \
        -command "set promptfmt \"$(${HOME}/.local/libexec/lf_prompt)\""

    dir="$(cat "$LF_DIRFILE")"
    rm "$LF_DIRFILE" &>/dev/null

    if [ -d "$dir" ]; then
        if [ "$dir" != "$(pwd)" ]; then
            cd "$dir"
        fi
    fi

    reset_prompt
    write_title_wd
    reset_beam
}
zle -N lf-wrap

function lazygit-wrap() {
    [ ! -d "${PWD}/.git" ] && [[ $(read -erk "?Not in a git repository. Create a new git repository? (y/n): ") =~ ^[Yy]$ ]] && git init "$PWD"
    [ -d "${PWD}/.git" ] && lazygit -p "$PWD"
    reset_prompt
}
zle -N lazygit-wrap

function workon_cwd() {
    local PROJECT_ROOT="${PWD}"

    while [ "$PROJECT_ROOT" != "/" ] && [ ! -e "$PROJECT_ROOT/.venv" ] && [ ! -d "$PROJECT_ROOT/.git" ]; do
        local PROJECT_ROOT="$(dirname "$PROJECT_ROOT")"
    done

    if [ -r "$PROJECT_ROOT/.venv" ]; then
        local ENV_NAME="$(cat $PROJECT_ROOT/.venv)"
    else
        local ENV_NAME="$(basename "$PROJECT_ROOT")"
    fi

    [ -n "$CD_VIRTUAL_ENV" ] && [ "$ENV_NAME" != "$CD_VIRTUAL_ENV" ] && {
        [ -n "$VIRTUAL_ENV" ] && deactivate
        unset CD_VIRTUAL_ENV
    }

    [ -n "$ENV_NAME" ] && [ ! "$(realpath "$VIRTUAL_ENV" 2>/dev/null)" = "$(realpath "$WORKON_HOME/$ENV_NAME")" ] && {
        if [ -e "$WORKON_HOME/$ENV_NAME/bin/activate" ]; then
            workon -n "$ENV_NAME" && export CD_VIRTUAL_ENV="$ENV_NAME"
        elif [ -e "$ENV_NAME/bin/activate" ]; then
            source "$ENV_NAME/bin/activate" && export CD_VIRTUAL_ENV="$ENV_NAME"
        else
            local ENV_NAME=""
        fi
    }

    [ -z "$ENV_NAME" ] && [ -n "$CD_VIRTUAL_ENV" ] && [ -n "$VIRTUAL_ENV" ] && {
        deactivate && unset CD_VIRTUAL_ENV
    }
}
add-zsh-hook chpwd workon_cwd

command -v direnv &>/dev/null && emulate zsh -c "$(direnv hook zsh)"