#!/bin/zsh

typeset -gA _zasync_fd=()       # slot → current tracked fd
typeset -gA _zasync_seq=()      # slot → sequence number of last-started job
typeset -gA _zasync_fd_slot=()  # fd   → slot name
typeset -gA _zasync_fd_seq=()   # fd   → sequence number at launch
typeset -gA _zasync_fd_pwd=()   # fd   → $PWD at launch
typeset -gA _zasync_fd_cb=()    # fd   → callback widget name
typeset -gA _zasync_reply=()    # slot → most-recent raw output
typeset -gi _zasync_n=0         # monotonic counter

zle -N .zasync.fd-callback

.zasync.start() {
  case ${1-} in
    :help-summary)
      print -r -- '  start <slot> <worker> <cb>    Start worker and register callback.'
      return 0
      ;;
    :help)
      print -r -- 'Usage: z-async start <slot> <worker> <cb>'
      print -r -- ''
      print -r -- 'Start worker and register callback.'
      print -r -- ''
      print -r -- 'Arguments:'
      print -r -- '  slot    Name used to track/cancel this in-flight job.'
      print -r -- '  worker  Command/function executed asynchronously.'
      print -r -- '  cb      ZLE widget called when worker output is ready.'
      return 0
      ;;
  esac

  (( $# == 3 )) || {
    .zasync.start :help
    return 2
  }

  local slot=$1 worker=$2 callback=$3

  # Cancel any previous in-flight job for this slot first.
  .zasync.cancel "$slot"

  local fd=
  sysopen -r -o cloexec -u fd <( "$worker" )

  local seq=$(( ++_zasync_n ))
  _zasync_fd[$slot]=$fd
  _zasync_seq[$slot]=$seq
  _zasync_fd_slot[$fd]=$slot
  _zasync_fd_seq[$fd]=$seq
  _zasync_fd_pwd[$fd]=$PWD
  _zasync_fd_cb[$fd]=$callback

  zle -Fw "$fd" .zasync.fd-callback
}

.zasync.cancel() {
  case ${1-} in
    :help-summary)
      print -r -- '  cancel <slot>                 Cancel in-flight worker for slot.'
      return 0
      ;;
    :help)
      print -r -- 'Usage: z-async cancel <slot>'
      print -r -- ''
      print -r -- 'Cancel in-flight worker for slot.'
      print -r -- ''
      print -r -- 'Arguments:'
      print -r -- '  slot    Name of the tracked slot to cancel.'
      return 0
      ;;
  esac

  (( $# == 1 )) || {
    .zasync.cancel :help
    return 2
  }

  local slot=$1
  local fd=${_zasync_fd[$slot]-}
  [[ -z $fd || $fd -lt 10 ]] && return 0

  unset "_zasync_fd[$slot]"
  unset "_zasync_fd_slot[$fd]" "_zasync_fd_seq[$fd]" "_zasync_fd_pwd[$fd]" "_zasync_fd_cb[$fd]"
  zle -F "$fd" 2>/dev/null
  exec {fd}<&-
}

.zasync.reply() {
  case ${1-} in
    :help-summary)
      print -r -- '  reply <slot>                  Print last reply for slot.'
      return 0
      ;;
    :help)
      print -r -- 'Usage: z-async reply <slot>'
      print -r -- ''
      print -r -- 'Print last reply for slot.'
      print -r -- ''
      print -r -- 'Arguments:'
      print -r -- '  slot    Name of the tracked slot.'
      return 0
      ;;
  esac

  (( $# == 1 )) || {
    .zasync.reply :help
    return 2
  }

  print -r -- "${_zasync_reply[$1]-}"
}

.zasync.help() {
  local cmd=${1-}

  if [[ -n $cmd ]]; then
    local func=.zasync.$cmd
    if [[ -v "functions[$func]" ]]; then
      "$func" :help
      return 0
    fi

    print -u2 -r -- "z-async: unknown command: $cmd"
    .zasync.help
    return 2
  fi

  print -r -- 'Usage: z-async help [command]'
  print -r -- 'Usage: z-async <command> [args]'
  print -r -- ''
  print -r -- 'Commands:'
  .zasync.start :help-summary
  .zasync.cancel :help-summary
  .zasync.reply :help-summary
}

.zasync.fd-callback() {
  local fd=$1
  zle -F "$fd"

  local slot=${_zasync_fd_slot[$fd]-}
  local req_seq=${_zasync_fd_seq[$fd]-0}
  local req_pwd=${_zasync_fd_pwd[$fd]-}
  local callback=${_zasync_fd_cb[$fd]-}

  local data=
  {
    # Avoid interrupting yank/kill widget sequences.
    if (( ${YANK_ACTIVE:-0} )); then
      local before=
      (( YANK_START <= CURSOR && YANK_END <= CURSOR )) && before=before
      builtin zle -f yank$before
    else
      builtin zle -f kill
    fi

    data=$( <&$fd )
  } always {
    exec {fd}<&-
    if [[ -n $slot ]]; then
      unset "_zasync_fd_slot[$fd]" "_zasync_fd_seq[$fd]" "_zasync_fd_pwd[$fd]" "_zasync_fd_cb[$fd]"
      if [[ -v _zasync_fd[$slot] && ${_zasync_fd[$slot]} == $fd ]]; then
        unset "_zasync_fd[$slot]"
      fi
    fi
  }

  # Discard stale results (wrong generation or wrong directory).
  local cur_seq=${_zasync_seq[$slot]-0}
  (( req_seq > 0 && req_seq == cur_seq )) && [[ $req_pwd == $PWD ]] || return 0

  _zasync_reply[$slot]=$data
  [[ -n $callback ]] && zle "$callback"
}

z-async() {
  (( $# > 0 )) || {
    .zasync.help
    return 2
  }

  if [[ $1 == help ]]; then
    shift
    .zasync.help "$@"
    return $?
  fi

  local func=.zasync.$1
  if [[ ! -v "functions[$func]" ]]; then
    print -u2 -r -- "z-async: unknown command: $1"
    .zasync.help
    return 2
  fi

  shift
  "$func" "$@"
}

z-async "$@"
