#compdef ruby irb erb ruby-mri -P (ruby|[ei]rb)[0-9.]# -value-,RUBY(LIB|OPT|PATH),-default-

local curcontext="$curcontext" state line expl desc RUBY variant ret=1
typeset -A opt_args
local -a opts irb erb all common charsets suf jit_opts

all=(
  '*-r+[require the library before executing your script]:library name:->library'
  '(--internal-encoding)-U[set internal encoding to UTF-8]'
  '(-E --encoding -K --external-encoding)'{-E+,--encoding=}'[specify default external and internal character encodings]:charset:->charsets'
  '(1 * -)'{-h,--help}'[print help message]'
  '(1 * -)--version[print the version]'
)

common=(
  '*-I+[specify $LOAD_PATH directory (may be used more than once)]:library directory:_files -/'
  '-w[turn warnings on for your script]'
  '-W-[set warning level]:warning level:((0\:silent 1\:medium 2\:verbose \\\:deprecated \\\:experimental))'
  '(-)1:ruby script:_files -g "*.rb(-.)"'
  '*::script argument:= ->normal'
)

opts=(
  '-0-[specify record separator]:input record separator in octal [\0]'
  '-a[autosplit mode with -n or -p (splits $_ into $F)]'
  '-c[check syntax only]'
  '-C+[cd to directory before executing your script]:directory:_files -/'
  '(-d --debug)'{-d,--debug}'[set debugging flags (set $DEBUG to true)]'
  "(1)*-e+[one line script (several -e's allowed, omit program file)]:one line of script"
  '-F-[split() pattern for autosplit (-a)]:input field separator:'
  '-i-[edit ARGV files in place (make backup if extension supplied)]:suffix for in-place-edit mode:(.bak)'
  '-K-[specifies KANJI (Japanese) code-set]:KANJI encoding:((e\:EUC-JP s\:Shift_JIS u\:UTF-8 N\:None))'
  "-l[assume 'while gets(); ... end' loop around your script]"
  "(-p)-n[assume 'while gets(); ... end' loop around your script]"
  '(-n)-p[assume loop like -n but print line also like sed]'
  '-s[enable some switch parsing for switches after script name]'
  '-S[look for the script using PATH environment variable]'
  '!-T-[turn on tainting checks]::taint level [1]:((0\:strings\ from\ streams/environment/ARGV\ are\ tainted 1\:no\ dangerous\ operation\ by\ tainted\ value 2\:process/file\ operations\ prohibited 3\:all\ generated\ objects\ are\ tainted 4\:no\ global\ \(non-tainted\)\ variable\ modification/no\ direct\ output))'
  '(--verbose)-v[print version number, then turn on verbose mode]'
  '(-v)--verbose[turn on verbose mode and disable script from stdin]'
  '-x-[strip off text before #!ruby line and perhaps cd to directory]:directory:_files -/'
  '(1 * -)--copyright[print the copyright]'
  --{en,dis}'able=[enable or disable features]:feature:(gems error_highlight did_you_mean syntax_suggest rubyopt frozen_string_literal mjit yjit all)'
  \!--{en,dis}able-{gems,rubyopt,all}
  '--dump=[dump debug information]:information:_sequence compadd - insns insns_without_opt yydebug parsetree parsetree_with_comment'
  --{external,internal}'-encoding=:charset:->charsets'
  '--backtrace-limit=[limit the maximum length of backtrace]:number'
  '!'{-y,--yydebug}
  '!--dump=:target:(version copyright usage yydebug syntax parsetree parsetree_with_comment insns)'
)

if [[ $service == ruby[0-9.]# ]]; then
  _pick_variant -r variant -c "$service --version" \
    ruby_no_jit='ruby ([01].|2.[0-5])' \
    ruby_experimental_jit='ruby (2.[67]|3.0)' \
    ruby31='ruby 3.1' ruby32='ruby 3.2' ruby33='ruby 3.3' ruby34='ruby 3.4' \
    ruby40='ruby 4.0' ruby41_higher

  if [[ $variant != "ruby_no_jit" ]]; then
    jit_opts=('(--jit --mjit --yjit --rjit --zjit)--jit[enable default JIT for the platform]')
  fi

  case "$variant" in
    (ruby_experimental_jit|ruby31)
      jit_opts+=(
        '--jit-warnings[enable printing JIT warnings]'
        '--jit-debug[enable JIT debugging (very slow)]'
        '--jit-wait[wait until JIT compilation finishes every time (for testing)]'
        '--jit-save-temps[save JIT temporary files]'
        '--jit-verbose=-[print JIT logs of level num or less to stderr]:log level'
        '--jit-max-cache=-[specify max number of methods to be JIT-ed in a cache]:num'
        '--jit-min-calls=-[specify number of calls to trigger JIT]:num'
      )
      ;|
    (ruby31|ruby32)
      jit_opts+=(
        '(--jit --mjit --yjit)--mjit[enable C compiler-based JIT compiler]'
        '(--jit --mjit --yjit)--yjit[enable in-process JIT compiler]'
      )
      ;|
    (ruby31|ruby32)
      jit_opts+=(
        '--mjit-warnings[enable printing JIT warnings]'
        '--mjit-debug[enable JIT debugging (very slow)]'
        '--mjit-wait[wait until JIT compilation finishes every time (for testing)]'
        '--mjit-save-temps[save JIT temporary files]'
        '--mjit-verbose=-[print JIT logs of level num or less to stderr]:log level'
        '--mjit-max-cache=-[specify max number of methods to be JIT-ed in a cache]:num'
        '--mjit-min-calls=-[specify number of calls to trigger JIT]:num'
        '--yjit-max-versions=[Maximum number of versions per basic block (default: 4)]:num'
        '--yjit-greedy-versioning[Greedy versioning mode (default: disabled)]'
      )
      ;|
    (ruby32|ruby33|ruby34|ruby40|ruby41_higher)
      jit_opts+=(
        '--yjit-stats[Enable collecting YJIT statistics]'
        '--yjit-exec-mem-size=[Hard limit on executable memory block in MiB]:mem size'
        '--yjit-call-threshold=[Number of calls to trigger JIT]:num'
      )
      ;|
    (ruby33|ruby34|ruby40|ruby41_higher)
      # mjit was removed since Ruby 3.3
      jit_opts+=(
        '--yjit-cold-threshold=[Global calls after which ISEQs not compiled (default: 200K)]:num'
        '--yjit-disable[Disable YJIT for lazily enabling it with RubyVM::YJIT.enable]'
        '--yjit-code-gc[Run code GC when the code size reaches the limit]'
        '--yjit-perf[Enable frame pointers and perf profiling]'
        '--yjit-trace-exits[Record Ruby source location when exiting from generated code]'
        '--yjit-trace-exits-sample-rate=[Trace exit locations only every Nth occurrence]:num'
      )
      ;|
    (ruby34|ruby40|ruby41_higher)
      jit_opts+=(
        '--yjit-mem-size=[Soft limit on YJIT memory usage in MiB (default: 128)]:mem size'
        "--yjit-log=[Enable logging of YJIT's compilation activity]: :_files"
      )
      ;|
    (ruby33|ruby34)
      # ruby 3.3 and 3.4 provide rjit
      jit_opts+=(
        '(--jit --yjit --rjit)--yjit[Enable in-process JIT compiler]'
        '--rjit-exec-mem-size=[Size of executable memory block in MiB (default: 64)]:mem size'
        '--rjit-call-threshold=[Number of calls to trigger JIT (default: 10)]:num'
        '--rjit-stats[Enable collecting RJIT statistics]'
        '--rjit-disable[Disable RJIT for lazily enabling it with RubyVM::RJIT.enable]'
        '--rjit-trace[Allow TracePoint during JIT compilation]'
        '--rjit-trace-exits[Trace side exit locations]'
      )
      ;|
    (ruby40|ruby41_higher)
      # zjit was introduced since ruby 4
      jit_opts+=(
        '(--jit --yjit --zjit)--yjit[Enable in-process JIT compiler]'
        '(--jit --yjit --zjit)--zjit[enable method-based JIT compiler]'
        '--zjit-mem-size=[Max amount of memory that ZJIT can use in MiB (default: 128)]:mem size'
        '--zjit-call-threshold=[Number of calls to trigger JIT (default: 30)]:num'
        '--zjit-num-profiles=[Number of profiled calls before JIT (default: 5)]:num'
        '--zjit-stats-quiet[Collect ZJIT stats and suppress output]'
        '--zjit-stats=-[Collect ZJIT stats]: :_files'
        '--zjit-disable[Disable ZJIT for lazily enabling it with RubyVM::ZJIT.enable]'
        '--zjit-perf[Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf]'
        '--zjit-log-compiled-iseqs=[Log compiled ISEQs to the file. The file will be truncated]: :_files'
        '--zjit-trace-exits=-[Record source on side-exit]:counter'
        '--zjit-trace-exits-sample-rate=[Frequency at which to record side exits]:frequency'
      )
      ;|
    (ruby41_higher)
      jit_opts+=(
        '--zjit-trace-compiles[Record compilation phases as Perfetto trace events]'
        '--zjit-trace-invalidation[Record invalidation events as Perfetto trace events]'
      )
      ;|
  esac
fi

irb=(
  '-f[suppress read of ~/.irbrc]'
  $opts[(r)*-d\[*]
  '(--noinspect)--inspect[use inspect for output]'
  "(--inspect)--noinspect[don't use inspect for output]"
  "--no-pager[don't use pager]"
  '(--prompt --prompt-mode --inf-ruby-mode --simple-prompt --noprompt)'{--prompt=,--prompt-mode=}'[switch prompt mode]:prompt mode:(default classic simple inf-ruby xmp null)'
  '(--prompt --prompt-mode --inf-ruby-mode --simple-prompt --noprompt)'{--inf-ruby-mode,--simple-prompt,--noprompt}
  '--tracer[display trace for each command execution]'
  '--back-trace-limit[set limit for backtraces]:limit [16]:'
  '!--irb_debug:level'
  '--context-mode:n'
  '--extra-doc-dir[add an extra doc dir for the doc dialog]'
  '(--noecho)--echo[show result]'
  "(--echo)--noecho[don't show result]"
  '(--noecho-on-assignment)--echo-on-assignment[show result on assignment]'
  "(--echo-on-assignment)--noecho-on-assignment[don't show result on assignment]"
  '--truncate-echo-on-assignment[show truncated result on assignment]'
  '(--nomultiline)--multiline[use multiline editor module]'
  "(--multiline)--nomultiline[don't use multiline editor module]"
  '(--nosingleline)--singleline[use single line editor module]'
  "(--singleline)--nosingleline[don't use single line editor module]"
  '(--nocolorize)--colorize[use color-highlighting]'
  "(--colorize)--nocolorize[don't use color-highlighting]"
  '(--noautocomplete)--autocomplete[use auto-completion]'
  "(--autocomplete)--noautocomplete[don't use auto-completion]"
  '(--regexp-completor)--type-completor[use regexp based completion]'
  '(--type-completor)--regexp-completor[use type based completion]'
  '(--noscript)--script[script mode]'
  '(--script)--noscript[no script mode]'
  '--single-irb[share self with sub-irb]'
  '(--noverbose)--verbose[show details]'
  "(--verbose)--noverbose[don't show details]"
)

erb=(
  "-P[don't evaluate lines which start with %]"
  '-T[specify trim mode]:mode [0]:((0\:EOL\ remains 1\:EOL\ removed\ if\ line\ ends\ with\ %\> 2\:EOL\ removed\ if\ line\ starts\ with\ \<%\ and\ ends\ with\ %\> -\:EOL\ is\ removed\ if\ line\ ends\ with\ -%\>,\ leading\ whitespace\ removed\ after\ \<%-))'
  '(-d --debug)'{-d,--debug}'[set debugging flags (set $DEBUG to true)]'
  '-n[print ruby script with line number]'
  '-x[print ruby script]'
  '-v[enable verbose mode]'
)

case "$service" in
  *RUBY(LIB|PATH)*)
    _dir_list && return 0
  ;;
  *RUBYOPT*)
    compset -q
    words=( fake "$words[@]" )
    (( CURRENT++ ))
    _arguments -C -s : "${(@M)opts:#*-[IdvwrK](|[-+])\[*}" && ret=0
    RUBY=ruby
  ;;
  ruby[0-9.]#)
    RUBY=$words[1]
    _arguments -C -s -S : $opts $all $common $jit_opts && ret=0
  ;;
  irb[0-9.]#)
    RUBY=${words[1]/irb/ruby}
    _arguments -C -s : $irb $all $common && ret=0
  ;;
  erb[0-9.]#)
    RUBY=${words[1]/erb/ruby}
    _arguments -C -s : $erb $all \
      '*:erb file:_files -g "*.erb(-.)"' && ret=0
  ;;
esac

case "$state" in
  normal)
    if (( ! $+opt_args[-e] )); then
      shift words
      (( CURRENT-- ))
    fi
    _normal && ret=0
  ;;
  charsets)
    charsets=( $(_call_program charsets $RUBY -e 'puts\ Encoding.list' 2>/dev/null) )
    # could also add Encoding.aliases.map(&:first) for aliases
    desc='charset'
    if [[ $curcontext = *option-(E|-encoding)-1 ]]; then
      if compset -P '*:'; then
	desc='internal charset'
      else
	desc='external charset'
	compset -S ':*' || suf=( -S: )
      fi
    fi
    _wanted charsets expl $desc compadd $suf -a charsets && ret=0
  ;;
  library)
    local -a dirs
    dirs=( $(_call_program libraries $RUBY -e 'puts\ \$:' 2>/dev/null) ${(s.:.)opt_args[-I]} )
    _wanted libraries expl library _path_files -W dirs && ret=0
  ;;
esac

return ret
