nighthawk 0.2.0

#!/usr/bin/env zsh

# nighthawk zsh plugin — inline ghost text autocomplete

#

# Install: add to ~/.zshrc:  source /path/to/nighthawk.zsh

# Requires: socat, jq


# --- Configuration ---

NIGHTHAWK_SOCKET="${NIGHTHAWK_SOCKET:-/tmp/nighthawk-$(id -u).sock}"
NIGHTHAWK_FUZZY_DISPLAY="${NIGHTHAWK_FUZZY_DISPLAY:-hint}"

# --- State ---

typeset -g _nh_suggestion=""

typeset -g _nh_replace_start=""

typeset -g _nh_replace_end=""

typeset -g _nh_last_buffer=""

typeset -g _nh_has_highlight=0

typeset -g _nh_diff_ops=""

typeset -g _nh_original_buffer=""

typeset -g _nh_original_cursor=""

typeset -g _nh_backoff_until=0


# --- Dependency check ---

if ! command -v socat &>/dev/null; then
    echo "nighthawk: socat not found, install with: apt install socat" >&2

    return 1

fi

if ! command -v jq &>/dev/null; then
    echo "nighthawk: jq not found, install with: apt install jq" >&2

    return 1

fi


# --- Ghost text rendering via POSTDISPLAY ---

_nh_render_ghost() {
    local ghost="$1"

    if [[ -n "$ghost" ]]; then
        POSTDISPLAY="$ghost"
        region_highlight+=("${#BUFFER} $((${#BUFFER} + ${#ghost})) fg=8")
        _nh_has_highlight=1
    fi

}

# --- Inline diff rendering for fuzzy matches ---

# Temporarily replaces the mistyped token in BUFFER with the diff text

# (keeps + deletes + inserts interleaved), then uses region_highlight

# to color each character: red+bold for deletes, gray for inserts.

# Original buffer is saved and restored on clear/accept.

_nh_render_diff() {
    local diff_ops="$1"

    local token_start="$2"

    local replace_end="$3"


    # Save original buffer so we can restore it later

    _nh_original_buffer="$BUFFER"
    _nh_original_cursor="$CURSOR"

    # Build the diff-rendered token and track highlight regions

    local new_token=""

    local -a diff_highlights=()

    local pos=$token_start

    local last_typed_pos=$token_start  # tracks end of user-typed chars


    local entry op ch

    for entry in ${(s: :)diff_ops}; do
        op="${entry[1]}"
        ch="${entry[3,-1]}"
        case "$op" in
            k)  new_token+="$ch"; pos=$((pos + 1))
                last_typed_pos=$pos ;;
            d)  new_token+="$ch"
                diff_highlights+=("$pos $((pos + 1)) fg=red,bold")
                pos=$((pos + 1))
                last_typed_pos=$pos ;;
            i)  new_token+="$ch"
                diff_highlights+=("$pos $((pos + 1)) fg=8")
                pos=$((pos + 1)) ;;
        esac
    done


    # Replace token in BUFFER: before + diff_token + after

    # Zsh strings are 1-indexed; replace_start/end are 0-indexed byte offsets

    local before="${_nh_original_buffer[1,$token_start]}"

    local after="${_nh_original_buffer[$((replace_end + 1)),-1]}"

    BUFFER="${before}${new_token}${after}"
    # Cursor after last Keep/Delete (user-typed chars), not trailing Inserts.

    # This way Insert chars after the cursor look like ghost text.

    CURSOR=$last_typed_pos

    # Update last_buffer so pre_redraw doesn't re-trigger on our modification

    _nh_last_buffer="$BUFFER"

    # Apply highlights

    for hl in "${diff_highlights[@]}"; do
        region_highlight+=("$hl")
    done

    _nh_has_highlight=${#diff_highlights[@]}
}

# --- Hint rendering for fuzzy matches ---

# Shows " -> suggestion" as gray POSTDISPLAY text.

# Does NOT modify BUFFER, so no save/restore dance needed.

_nh_render_hint() {
    local suggestion="$1"

    if [[ -n "$suggestion" ]]; then
        POSTDISPLAY=" → $suggestion"
        region_highlight+=("${#BUFFER} $((${#BUFFER} + ${#POSTDISPLAY})) fg=8")
        _nh_has_highlight=1
    fi

}

# --- Restore original buffer before user edits ---

# When diff rendering has modified BUFFER, we must restore the original

# before any editing operation so the keystroke applies to the right text.

_nh_restore_before_edit() {
    if [[ -n "$_nh_original_buffer" ]]; then
        BUFFER="$_nh_original_buffer"
        CURSOR=${#BUFFER}  # end of original text

        _nh_original_buffer=""
        _nh_original_cursor=""
        region_highlight=()
        _nh_has_highlight=0
        _nh_diff_ops=""
        _nh_suggestion=""
        _nh_replace_start=""
        _nh_replace_end=""
        unset POSTDISPLAY

        _nh_last_buffer="$BUFFER"
    fi

}

_nh_self_insert() {
    _nh_restore_before_edit

    zle .self-insert

}
zle -N self-insert _nh_self_insert


_nh_backward_delete() {
    _nh_restore_before_edit

    zle .backward-delete-char

}
zle -N backward-delete-char _nh_backward_delete


_nh_clear_ghost() {
    unset POSTDISPLAY


    # Remove highlight entries we added

    if (( _nh_has_highlight )); then
        local i

        for i in $(seq $_nh_has_highlight); do
            region_highlight[-1]=()
        done

        _nh_has_highlight=0
    fi


    # Restore original buffer if diff rendering modified it

    if [[ -n "$_nh_original_buffer" ]]; then
        BUFFER="$_nh_original_buffer"
        CURSOR="$_nh_original_cursor"
        _nh_original_buffer=""
        _nh_original_cursor=""
    fi


    _nh_suggestion=""
    _nh_replace_start=""
    _nh_replace_end=""
    _nh_diff_ops=""
}

# --- Auto-start ---

_nh_ensure_daemon() {
    # Quick check: does the socket exist?

    [[ -S "$NIGHTHAWK_SOCKET" ]] && return 0


    # Try to start the daemon (one attempt, non-blocking)

    if command -v nh &>/dev/null; then
        nh start &>/dev/null

    elif command -v nighthawk-daemon &>/dev/null; then
        nighthawk-daemon &>/dev/null &
        disown

        sleep 0.2

    fi


    return 0

}

# --- Daemon communication ---

_nh_query() {
    # Back off for 5s after connection failure so dead daemon doesn't block every keystroke

    local now=$EPOCHSECONDS

    if (( now < _nh_backoff_until )); then
        return

    fi


    _nh_ensure_daemon


    local buffer="$1"

    local cursor="$2"


    # Escape for JSON: backslashes then double quotes

    local escaped_buffer="${buffer//\\/\\\\}"

    escaped_buffer="${escaped_buffer//\"/\\\"}"
    local escaped_cwd="${PWD//\\/\\\\}"

    escaped_cwd="${escaped_cwd//\"/\\\"}"

    local json="{\"input\":\"${escaped_buffer}\",\"cursor\":${cursor},\"cwd\":\"${escaped_cwd}\",\"shell\":\"zsh\"}"


    local response

    response=$(echo "$json" | socat -t1 - UNIX-CONNECT:"$NIGHTHAWK_SOCKET" 2>/dev/null)

    if [[ -z "$response" ]]; then
        _nh_backoff_until=$(( EPOCHSECONDS + 5 ))
        return

    fi


    # Parse first suggestion (text, replace range, and diff_ops if present)

    local text replace_start replace_end diff_ops_str

    eval $(echo "$response" | jq -r '
        if (.suggestions | length) > 0 then
            "text=" + (.suggestions[0].text | @sh) +
            " replace_start=" + (.suggestions[0].replace_start | tostring) +
            " replace_end=" + (.suggestions[0].replace_end | tostring) +
            " diff_ops_str=" + ((.suggestions[0].diff_ops // null) | if . then [.[] | .op[0:1] + ":" + .ch] | join(" ") | @sh else ("" | @sh) end)
        else
            "text='"''"'"
        end
    ' 2>/dev/null)


    if [[ -n "$text" ]]; then
        _nh_suggestion="$text"
        _nh_replace_start="$replace_start"
        _nh_replace_end="$replace_end"

        if [[ -n "$diff_ops_str" ]]; then
            # Fuzzy match: render based on display mode

            if [[ "$NIGHTHAWK_FUZZY_DISPLAY" == "hint" ]]; then
                _nh_render_hint "$text"

            else
                _nh_diff_ops="$diff_ops_str"
                _nh_render_diff "$diff_ops_str" "$replace_start" "$replace_end"

            fi

        else
            # Check if this is a true prefix match or a replacement

            local already_typed_len=$(( cursor - replace_start ))

            if (( already_typed_len >= 0 && already_typed_len < ${#text} )); then
                local typed_part="${buffer:$replace_start:$already_typed_len}"

                if [[ "${text:0:$already_typed_len}" == "$typed_part" ]]; then
                    # True prefix match: show suffix as ghost text

                    local ghost="${text:$already_typed_len}"

                    _nh_render_ghost "$ghost"

                else
                    # Replacement changes typed text: show hint instead

                    _nh_render_hint "$text"

                fi

            fi

        fi

    fi

}

# --- ZLE hooks ---


# Save existing hook if any

if zle -l zle-line-pre-redraw; then
    zle -A zle-line-pre-redraw _nh_orig_pre_redraw

fi


_nh_pre_redraw() {
    # Chain to original hook

    (( $+functions[_nh_orig_pre_redraw] )) && _nh_orig_pre_redraw


    # Only act if buffer changed

    [[ "$BUFFER" == "$_nh_last_buffer" ]] && return

    _nh_last_buffer="$BUFFER"

    # Clear previous ghost text

    _nh_clear_ghost


    # Only suggest when cursor is at end of line

    (( CURSOR != ${#BUFFER} )) && return


    # Need at least 2 chars

    (( ${#BUFFER} < 2 )) && return


    _nh_query "$BUFFER" "$CURSOR"

}

zle -N zle-line-pre-redraw _nh_pre_redraw


# --- Accept suggestion ---

_nh_accept() {
    if [[ -n "$_nh_suggestion" && "$_nh_replace_start" != "" && "$_nh_replace_end" != "" ]]; then
        local suggestion="$_nh_suggestion"

        local rstart="$_nh_replace_start"

        local rend="$_nh_replace_end"


        # If diff rendering modified BUFFER, restore original first

        # so replace_start/end offsets are correct

        if [[ -n "$_nh_original_buffer" ]]; then
            BUFFER="$_nh_original_buffer"
            _nh_original_buffer=""
            _nh_original_cursor=""
        fi


        # Clear ghost/highlight state

        unset POSTDISPLAY

        region_highlight=()
        _nh_has_highlight=0
        _nh_suggestion=""
        _nh_diff_ops=""

        # Replace the token: BUFFER[0..rstart] + suggestion + BUFFER[rend..]

        # Zsh strings are 1-indexed

        local before="${BUFFER[1,$rstart]}"

        local after="${BUFFER[$((rend + 1)),-1]}"

        BUFFER="${before}${suggestion}${after}"
        CURSOR=${#BUFFER}

        _nh_replace_start=""
        _nh_replace_end=""
        _nh_last_buffer="$BUFFER"
        zle redisplay

    else
        # No suggestion — fall through to default Tab

        zle expand-or-complete

    fi

}

zle -N _nh_accept

bindkey '^I' _nh_accept


# --- Clean up on line accept (Enter) ---

_nh_line_finish() {
    # Restore original buffer if diff rendering modified it,

    # so the shell executes what the user actually typed

    if [[ -n "$_nh_original_buffer" ]]; then
        BUFFER="$_nh_original_buffer"
        _nh_original_buffer=""
        _nh_original_cursor=""
    fi


    unset POSTDISPLAY

    region_highlight=()
    _nh_has_highlight=0
    _nh_suggestion=""
    _nh_replace_start=""
    _nh_replace_end=""
    _nh_diff_ops=""
    _nh_last_buffer="$BUFFER"
    zle accept-line

}

zle -N _nh_line_finish

bindkey '^M' _nh_line_finish