#!/bin/zsh
#
# Enables integration between zsh and kitty based on KITTY_SHELL_INTEGRATION.
# The latter is set by kitty based on kitty.conf.
#
# This is an autoloadable function. It's invoked automatically in shells
# directly spawned by kitty but not in any other shells. For example, running
# `exec zsh`, `sudo -E zsh`, `tmux`, or plain `zsh` will create a shell where
# kitty-integration won't automatically run. Zsh users who want integration with
# kitty in all shells should add the following lines to their .zshrc:
#
#   if [[ -n "$KITTY_INSTALLATION_DIR" ]]; then
#     export KITTY_SHELL_INTEGRATION="enabled"
#     autoload -Uz -- "$KITTY_INSTALLATION_DIR"/shell-integration/zsh/kitty-integration
#     kitty-integration
#     unfunction kitty-integration
#   fi
#
# Implementation note: We can assume that alias expansion is disabled in this
# file, so no need to quote defensively. We still have to defensively prefix all
# builtins with `builtin` to avoid accidentally invoking user-defined functions.
# We avoid `function` reserved word as an additional defensive measure.

builtin emulate -L zsh -o no_warn_create_global -o no_aliases

[[ -o interactive ]]                || builtin return 0  # non-interactive shell
[[ -n "$KITTY_SHELL_INTEGRATION" ]] || builtin return 0  # integration disabled
(( ! $+_ksi_state ))                || builtin return 0  # already initialized

# 0: no OSC 133 [AC] marks have been written yet.
# 1: the last written OSC 133 C has not been closed with D yet.
# 2: none of the above.
builtin typeset -gi _ksi_state

# Attempt to create a writable file descriptor to the TTY so that we can print
# to the TTY later even when STDOUT is redirected. This code is fairly subtle.
#
# - It's tempting to do `[[ -t 1 ]] && exec {_ksi_state}>&1` but we cannot do this
#   because it'll create a file descriptor >= 10 without O_CLOEXEC. This file
#   descriptor will leak to child processes.
# - If we do `exec {3}>&1`, the file descriptor won't leak to the child processes
#   but it'll still leak if the current process is replaced with another. In
#   addition, it'll break user code that relies on fd 3 being available.
# - Zsh doesn't expose dup3, which would have allowed us to copy STDOUT with
#   O_CLOEXEC. The only way to create a file descriptor with O_CLOEXEC is via
#   sysopen.
# - `zmodload zsh/system` and `sysopen -o cloexec -wu _ksi_fd -- /dev/tty` can
#   fail with an error message to STDERR (the latter can happen even if /dev/tty
#   is writable), hence the redirection of STDERR. We do it for the whole block
#   for performance reasons (redirections are slow).
# - We must open the file descriptor right here rather than in _ksi_deferred_init
#   because there are broken zsh plugins out there that run `exec {fd}< <(cmd)`
#   and then close the file descriptor more than once while suppressing errors.
#   This could end up closing our file descriptor if we opened it in
#   _ksi_deferred_init.
typeset -gi _ksi_fd
{
    builtin zmodload zsh/system && (( $+builtins[sysopen] )) && {
        { [[ -w     $TTY ]] && builtin sysopen -o cloexec -wu _ksi_fd --     $TTY } ||
        { [[ -w /dev/tty ]] && builtin sysopen -o cloexec -wu _ksi_fd -- /dev/tty }
    }
} 2>/dev/null || (( _ksi_fd = 1 ))

# Asks kitty to print $@ to its STDERR. This is for debugging.
_ksi_debug_print() {
    builtin local data
    data=$(builtin command base64 <<<"${(j: :)@}") || builtin return
    # Removing all spaces rather than just \n allows this code to
    # work on broken systems where base64 outputs \r\n.
    builtin print -nu "$_ksi_fd" '\eP@kitty-print|'"${data//[[:space:]]}"'\e\\'
}

# We defer initialization until precmd for several reasons:
#
# - Oh My Zsh and many other configs remove zle-line-init and
#   zle-line-finish hooks when they initialize.
# - By deferring initialization we allow user rc files to opt out from some
#   parts of integration. For example, if a zshrc theme prints OSC 133
#   marks, it can append " no-prompt-mark" to KITTY_SHELL_INTEGRATION during
#   initialization to avoid redundant marks from our code.
builtin typeset -ag precmd_functions
precmd_functions+=(_ksi_deferred_init)

_ksi_deferred_init() {
    builtin emulate -L zsh -o no_warn_create_global -o no_aliases

    # Recognized options: no-cursor, no-title, no-prompt-mark, no-complete, no-cwd, no-sudo.
    builtin local -a opt
    opt=(${(s: :)KITTY_SHELL_INTEGRATION})
    builtin unset KITTY_SHELL_INTEGRATION

    if [[ -n "$SSH_KITTEN_KITTY_DIR" ]]; then
        if [[ ! "$PATH" =~ (^|:)${SSH_KITTEN_KITTY_DIR}(:|$) ]] && [[ -z "$(builtin command -v kitten)" ]]; then
            builtin export PATH="${PATH}:${SSH_KITTEN_KITTY_DIR}"
        fi
        builtin unset SSH_KITTEN_KITTY_DIR
    fi

    # The directory where kitty-integration is located: /.../shell-integration/zsh.
    builtin local self_dir="${functions_source[_ksi_deferred_init]:A:h}"
    # The directory with _kitty. We store it in a directory of its own rather than
    # in $self_dir because we are adding it to fpath and we don't want any other
    # files to be accidentally autoloadable.
    builtin local comp_dir="$self_dir/completions"

    # Enable completions for `kitty` command.
    if (( ! opt[(Ie)no-complete] )) && [[ -r $comp_dir/_kitty ]]; then
        if (( $+functions[compdef] )); then
            # If compdef is defined, then either compinit has already run or it's
            # a shim that records all calls for the purpose of replaying them after
            # compinit. Either way we clobber the existing completion for kitty and
            # install our own.
            builtin unset "functions[_kitty]"
            builtin autoload -Uz -- $comp_dir/_kitty
            compdef _kitty kitty
            compdef _kitty clone-in-kitty
            compdef _kitty kitten
        fi

        # If compdef is not set, compinit has not run yet. In this case we must
        # add our completions directory to fpath so that _kitty gets picked up by
        # compinit.
        #
        # We extend fpath even if compinit has run because it might run again.
        # Without our completions directory in fpath compinit would our _comp
        # mapping.
        builtin typeset -ga fpath
        fpath=($comp_dir ${fpath:#$comp_dir})
    fi

    # Enable semantic markup with OSC 133.
    if (( ! opt[(Ie)no-prompt-mark] )); then
        _ksi_precmd() {
            builtin local -i cmd_status=$?
            builtin emulate -L zsh -o no_warn_create_global -o no_aliases

            # Don't write OSC 133 D when our precmd handler is invoked from zle.
            # Some plugins do that to update prompt on cd.
            if ! builtin zle; then
                # This code works incorrectly in the presence of a precmd or chpwd
                # hook that prints. For example, sindresorhus/pure prints an empty
                # line on precmd and marlonrichert/zsh-snap prints $PWD on chpwd.
                # We'll end up writing our OSC 133 D mark too late.
                #
                # Another failure mode is when the output of a command doesn't end
                # with LF and prompst_sp is set (it is by default). In this case
                # we'll incorrectly state that '%' from prompt_sp is a part of the
                # command's output.
                if (( _ksi_state == 1 )); then
                    # The last written OSC 133 C has not been closed with D yet.
                    # Close it and supply status.
                    builtin print -nu $_ksi_fd '\e]133;D;'$cmd_status'\a'
                    (( _ksi_state = 2 ))
                elif (( _ksi_state == 2 )); then
                    # There might be an unclosed OSC 133 C. Close that.
                    builtin print -nu $_ksi_fd '\e]133;D\a'
                fi
            fi

            builtin local mark1=$'%{\e]133;A\a%}'
            if [[ -o prompt_percent ]]; then
                builtin typeset -g precmd_functions
                if [[ ${precmd_functions[-1]} == _ksi_precmd ]]; then
                    # This is the best case for us: we can add our marks to PS1 and
                    # PS2. This way our marks will be printed whenever zsh
                    # redisplays prompt: on reset-prompt, on SIGWINCH, and on
                    # SIGCHLD if notify is set. Themes that update prompt
                    # asynchronously from a `zle -F` handler might still remove our
                    # marks. Oh well.
                    builtin local mark2=$'%{\e]133;A;k=s\a%}'
                    # Add marks conditionally to avoid a situation where we have
                    # several marks in place. These conditions can have false
                    # positives and false negatives though.
                    #
                    # - False positive (with prompt_percent): PS1="%(?.$mark1.)"
                    # - False negative (with prompt_subst):   PS1='$mark1'
                    [[ $PS1 == *$mark1* ]] || PS1=${mark1}${PS1}
                    # PS2 mark is needed when clearing the prompt on resize
                    [[ $PS2 == *$mark2* ]] || PS2=${mark2}${PS2}
                    (( _ksi_state = 2 ))
                else
                    # If our precmd hook is not the last, we cannot rely on prompt
                    # changes to stick, so we don't even try. At least we can move
                    # our hook to the end to have better luck next time. If there is
                    # another piece of code that wants to take this privileged
                    # position, this won't work well. We'll break them as much as
                    # they are breaking us.
                    precmd_functions=(${precmd_functions:#_ksi_precmd} _ksi_precmd)
                    # Plugins that invoke precmd hooks from zle do that before zle
                    # is trashed. This means that the cursor is in the middle of
                    # BUFFER and we cannot print our mark there. Prompt might
                    # already have a mark, so the following reset-prompt will write
                    # it. If it doesn't, there is nothing we can do.
                    if ! builtin zle; then
                        builtin print -rnu $_ksi_fd -- $mark1[3,-3]
                        (( _ksi_state = 2 ))
                    fi
                fi
            elif ! builtin zle; then
                # Without prompt_percent we cannot patch prompt. Just print the
                # mark, except when we are invoked from zle. In the latter case we
                # cannot do anything.
                builtin print -rnu $_ksi_fd -- $mark1[3,-3]
                (( _ksi_state = 2 ))
            fi
        }

        _ksi_preexec() {
            builtin emulate -L zsh -o no_warn_create_global -o no_aliases

            # This can potentially break user prompt. Oh well. The robustness of
            # this code can be improved in the case prompt_subst is set because
            # it'll allow us distinguish (not perfectly but close enough) between
            # our own prompt, user prompt, and our own prompt with user additions on
            # top. We cannot force prompt_subst on the user though, so we would
            # still need this code for the no_prompt_subst case.
            PS1=${PS1//$'%{\e]133;A\a%}'}
            PS2=${PS2//$'%{\e]133;A;k=s\a%}'}

            # This will work incorrectly in the presence of a preexec hook that
            # prints. For example, if MichaelAquilina/zsh-you-should-use installs
            # its preexec hook before us, we'll incorrectly mark its output as
            # belonging to the command (as if the user typed it into zle) rather
            # than command output.
            builtin print -nu "$_ksi_fd" -f '\e]133;C;cmdline=%q\a' "$1"
            (( _ksi_state = 1 ))
        }

        # the following two lines are commented out as currently kitty doesn't use B prompt marking
        # and hooking zle widgets in ZSH is a total minefield, see https://github.com/kovidgoyal/kitty/issues/4428
        # so we can at least tell users to use no-cursor and with that avoid hooking ZLE widgets at all
        # functions[_ksi_zle_line_init]+='
        #     builtin print -nu "$_ksi_fd" "\\e]133;B\\a"'
    fi

    # Enable reporting current working dir to terminal
    if (( ! opt[(Ie)no-cwd] )); then
        _ksi_report_pwd() { builtin print -nu $_ksi_fd '\e]7;kitty-shell-cwd://'"$HOST""$PWD"'\a'; }
        chpwd_functions=(${chpwd_functions[@]} "_ksi_report_pwd")
        # An executed program could change cwd and report the changed cwd, so also report cwd at each new prompt
        # as in this case chpwd_functions is insufficient. chpwd_functions is still needed for things like: cd x && something
        functions[_ksi_precmd]+="
            _ksi_report_pwd"
        _ksi_report_pwd
    fi

    # Enable terminal title changes.
    if (( ! opt[(Ie)no-title] )); then
        # We don't use `print -P` because it depends on prompt options, which
        # we don't control and cannot change.
        #
        # We use (V) in preexec to convert control characters to something visible
        # (LF becomes \n, etc.). This isn't necessary in precmd because (%) does it
        # for us.
        builtin local is_ssh_session="n"
        if [[ -n "$KITTY_PID" ]]; then
            # kitty running locally
        elif [[ -n "$SSH_TTY" || -n "$SSH2_TTY$KITTY_WINDOW_ID" ]]; then
            # connected to most SSH servers
            # or use ssh kitten to connected to some SSH servers that do not set SSH_TTY
            is_ssh_session="y"
        elif [[ -n "$(builtin command -v who)" ]]; then
            # the shell integration script is installed manually on the remote system
            # the environment variables are cleared after sudo
            # OpenSSH's sshd creates entries in utmp for every login so use those
            [[ "$(builtin command who -m 2> /dev/null)" =~ "\([a-fA-F.:0-9]+\)$" ]] && is_ssh_session="y"
        fi

        if [[ "$is_ssh_session" == "y" ]]; then
            functions[_ksi_precmd]+="
                builtin print -rnu $_ksi_fd \$'\\e]2;'\"\${(V)\${HOST-}/%.*/}: \${(%):-%(4~|…/%3~|%~)}\"\$'\\a'"
            functions[_ksi_preexec]+="
                builtin print -rnu $_ksi_fd \$'\\e]2;'\"\${(V)\${HOST-}/%.*/}: \${(V)1}\"\$'\\a'"
        else
            functions[_ksi_precmd]+="
                builtin print -rnu $_ksi_fd \$'\\e]2;'\"\${(%):-%(4~|…/%3~|%~)}\"\$'\\a'"
            functions[_ksi_preexec]+="
                builtin print -rnu $_ksi_fd \$'\\e]2;'\"\${(V)1}\"\$'\\a'"
        fi
    fi

    # Enable cursor shape changes depending on the current keymap.
    if (( ! opt[(Ie)no-cursor] )); then
        # This implementation leaks blinking block cursor into external commands
        # executed from zle. For example, users of fzf-based widgets may find
        # themselves with a blinking block cursor within fzf.
        _ksi_zle_line_init _ksi_zle_line_finish _ksi_zle_keymap_select() {
            case ${KEYMAP-} in
                # Blinking block cursor.
                vicmd|visual) builtin print -nu "$_ksi_fd" '\e[1 q';;
                # Blinking bar cursor.
                *)            builtin print -nu "$_ksi_fd" '\e[5 q';;
            esac
        }
        # Restore the blinking default shape before executing an external command
        functions[_ksi_preexec]+="
            builtin print -rnu $_ksi_fd \$'\\e[0 q'"
    fi


    # Some zsh users manually run `source ~/.zshrc` in order to apply rc file
    # changes to the current shell. This is a terrible practice that breaks many
    # things, including our shell integration. For example, Oh My Zsh and Prezto
    # (both very popular among zsh users) will remove zle-line-init and
    # zle-line-finish hooks if .zshrc is manually sourced. Prezto will also remove
    # zle-keymap-select.
    #
    # Another common (and much more robust) way to apply rc file changes to the
    # current shell is `exec zsh`. This will remove our integration from the shell
    # unless it's explicitly invoked from .zshrc. This is not an issue with
    # `exec zsh` but rather with our implementation of automatic shell integration.

    # In the ideal world we would use add-zle-hook-widget to hook zle-line-init
    # and similar widget. This breaks user configs though, so we have do this
    # horrible thing instead.
    builtin local hook func widget orig_widget flag
    for hook in line-init line-finish keymap-select; do
        func=_ksi_zle_${hook/-/_}
        (( $+functions[$func] )) || builtin continue
        widget=zle-$hook
        if [[ $widgets[$widget] == user:azhw:* &&
              $+functions[add-zle-hook-widget] -eq 1 ]]; then
            # If the widget is already hooked by add-zle-hook-widget at the top
            # level, add our hook at the end. We MUST do it this way. We cannot
            # just wrap the widget ourselves in this case because it would
            # trigger bugs in add-zle-hook-widget.
            add-zle-hook-widget $hook $func
        else
            if (( $+widgets[$widget] )); then
                # There is a widget but it's not from add-zle-hook-widget. We
                # can rename the original widget, install our own and invoke
                # the original when we are called.
                #
                # Note: The leading dot is to work around bugs in
                # zsh-syntax-highlighting.
                orig_widget=._ksi_orig_$widget
                builtin zle -A $widget $orig_widget
                if [[ $widgets[$widget] == user:* ]]; then
                    # No -w here to preserve $WIDGET within the original widget.
                    flag=
                else
                    flag=w
                fi
                functions[$func]+="
                    builtin zle $orig_widget -N$flag -- \"\$@\""
            fi
            builtin zle -N $widget $func
        fi
    done

    if (( $+functions[_ksi_preexec] )); then
        builtin typeset -ag preexec_functions
        preexec_functions+=(_ksi_preexec)
    fi

    builtin typeset -ag precmd_functions
    if (( $+functions[_ksi_precmd] )); then
        precmd_functions=(${precmd_functions:/_ksi_deferred_init/_ksi_precmd})
        _ksi_precmd
    else
        precmd_functions=(${precmd_functions:#_ksi_deferred_init})
    fi

    if [ -n "${KITTY_IS_CLONE_LAUNCH}" ]; then
        builtin local orig_conda_env="$CONDA_DEFAULT_ENV"
        builtin eval "${KITTY_IS_CLONE_LAUNCH}"
        builtin hash -r 2> /dev/null 1> /dev/null
        builtin local venv="${VIRTUAL_ENV}/bin/activate"
        builtin local sourced=""
        _ksi_s_is_ok() {
            [[ -z "$sourced" && "$KITTY_CLONE_SOURCE_STRATEGIES" == *",$1,"* ]] && builtin return 0
            builtin return 1
        }

        if _ksi_s_is_ok "venv" && [[ -n "${VIRTUAL_ENV}" && -r "$venv" ]]; then
            sourced="y"
            builtin unset VIRTUAL_ENV
            builtin source "$venv"
        fi; if _ksi_s_is_ok "conda" && [[ -n "${CONDA_DEFAULT_ENV}" && (( $+commands[conda] )) && "${CONDA_DEFAULT_ENV}" != "$orig_conda_env" ]]; then
            sourced="y"
            conda activate "${CONDA_DEFAULT_ENV}"
        fi; if _ksi_s_is_ok "env_var" && [[ -n "${KITTY_CLONE_SOURCE_CODE}" ]]; then
            sourced="y"
            builtin eval "${KITTY_CLONE_SOURCE_CODE}"
        fi; if _ksi_s_is_ok "path" && [[ -r "${KITTY_CLONE_SOURCE_PATH}" ]]; then
            sourced="y"
            builtin source "${KITTY_CLONE_SOURCE_PATH}"
        fi
        builtin unfunction _ksi_s_is_ok
        # Ensure PATH has no duplicate entries
        builtin typeset -gxU PATH="$PATH"
    fi
    builtin unset KITTY_IS_CLONE_LAUNCH KITTY_CLONE_SOURCE_STRATEGIES

    builtin alias edit-in-kitty="kitten edit-in-kitty"

    if (( ! opt[(Ie)no-sudo] )) ; then
        if [[ -n "$TERMINFO" && ! ( -r "/usr/share/terminfo/x/xterm-kitty" || -r "/usr/share/terminfo/78/xterm-kitty" ) ]]; then
            sudo() {
                # Ensure terminfo is available in sudo
                builtin local is_sudoedit="n"
                for arg; do
                    if [[ "$arg" == "-e" || $arg == "--edit" ]]; then
                        is_sudoedit="y"
                        builtin break;
                    fi
                    [[ "$arg" != -* && "$arg" != *=* ]] && builtin break  # command found
                done
                if [[ "$is_sudoedit" == "y" ]]; then
                    builtin command sudo "$@";
                else
                    builtin command sudo TERMINFO="$TERMINFO" "$@";
                fi
            }
        fi
    fi

    # Map alt+left/right to move by word if not already mapped. This is expected behavior on macOS and I am tired
    # of answering questions about it.
    [[ $(builtin bindkey "^[[1;3C") == *" undefined-key" ]] && builtin bindkey "^[[1;3C" "forward-word"
    [[ $(builtin bindkey "^[[1;3D") == *" undefined-key" ]] && builtin bindkey "^[[1;3D" "backward-word"

    # Unfunction _ksi_deferred_init to save memory. Don't unfunction
    # kitty-integration though because decent public functions aren't supposed to
    # to unfunction themselves when invoked. Unfunctioning is done by calling code.
    builtin unfunction _ksi_deferred_init
}

_ksi_transmit_data() {
    builtin local data="${1//[[:space:]]}"
    builtin local pos=0
    builtin local chunk_num=0
    while [ $pos -lt ${#data} ]; do
        builtin local chunk="${data:$pos:2048}"
        pos=$(($pos+2048))
        builtin print -nu "$_ksi_fd" -f '\eP@kitty-%s|%s:%s\e\\' "${2}" "${chunk_num}" "${chunk}"
        chunk_num=$(($chunk_num+1))
    done
    # save history so it is available in new shell
    [ "$3" = "save_history" ] && builtin fc -AI
    builtin print -nu "$_ksi_fd" -f '\eP@kitty-%s|\e\\' "${2}"
}

clone-in-kitty() {
    builtin local data="shell=zsh,pid=$$,cwd=$(builtin printf "%s" "$PWD" | builtin command base64)"
    while :; do
        case "$1" in
            "") break;;
            -h|--help)
                builtin printf "%s\n\n%s\n" "Clone the current zsh session into a new kitty window." "For usage instructions see: https://sw.kovidgoyal.net/kitty/shell-integration/#clone-shell"
                builtin return
                ;;
            *) data="$data,a=$(builtin printf "%s" "$1" | builtin command base64)";;
        esac
        shift
    done
    builtin local env
    builtin local env_vars
    builtin local varname
    env_vars=(${(f)"$(builtin export)"})
    for i in $env_vars; do
        varname="${i%%=*}"
        env="${env}$(builtin printf "%s=%s\0" "$varname" "${(P)varname}")"
    done
    data="$data,env=$(builtin printf "%s" "$env" | builtin command base64)"
    _ksi_transmit_data "$data" "clone" "save_history"
}