#compdef ceccomp
# Zsh completion script for ceccomp
local -a subcmd options expl arch_compstr output_compstr
local state

local archs=(
    'x86_64' 'i386' 'x32' 'aarch64' 'arm' 'loongarch64' 'm68k' 'mips' 'mipsel'
    'mips64' 'mipsel64' 'mips64n32' 'mipsel64n32' 'parisc' 'parisc64' 'ppc64' 'ppc64le'
    'ppc' 's390x' 's390' 'riscv64'
)
arch_compstr=(
    '(-a --arch)'{-a,--arch}"[Target BPF architecture]:ARCH:($archs)"
)
color_compstr=(
    '(-c --color)'{-c,--color}'[When to display in color]:WHEN:(always auto never)'
)
output_compstr=(
    '(-o --output)'{-o,--output}'[Print to file to avoid mixing tracee output]:FILE:_files'
)
local syscalls=(
    'open' 'read' 'write' 'close' 'mmap' 'mprotect' 'execve' 'execveat' 'pread64'
    'readv' 'writev' 'preadv' 'preadv2' 'openat' 'openat2' 'sendfile64' 'send'
    'sendto' 'sendmsg' 'recv' 'recvfrom' 'recvmsg' 'io_uring_setup' 'io_uring_enter'
    'io_uring_register' 'ptrace'
)

if (( CURRENT == 2 )) {
    subcmd=(
        'asm:Assemble bpf text to raw bytes'
        'disasm:Disassemble raw bytes to bpf text'
        'trace:Run program or trace pid, extract bpf filter and then print to text'
        'emu:Emulate bpf program with given syscall and bpf text'
        'probe:Trace the program for the first filter and emulate common syscalls'
        'help:Display ceccomp help information'
        'version:Display ceccomp version'
    )
    _describe 'subcmd' subcmd
    return
} elif (( CURRENT > 2 )) {
    case $words[2] {
        (asm)
            _arguments \
                $arch_compstr \
                $color_compstr \
                '(-f --fmt)'{-f,--fmt}'[Output format of BPF]:FMT:(raw hexline hexfmt)' \
                '2:BPF:_files' \
                '*: :'
            ;;
        (disasm)
            _arguments \
                $arch_compstr \
                $color_compstr \
                '2:RAW:_files' \
                '*: :'
            ;;
        (trace)
            _arguments -C \
                $arch_compstr \
                $color_compstr \
                '(-p --pid)'{-p,--pid}'[Attach to which process to extract its filters]:PID:->getpid' \
                $output_compstr \
                '*:arguments:_files'
            # complete non-kernel pids like kill
            if [[ $state == 'getpid' ]] {
                local line pids lines
                pids=()
                # extract command output line by line
                lines=("${(@f)$(ps --ppid 2 -p 2 -N -o pid=,tty=,user=,comm=)}")
                for line ($lines) {
                    pids+=${line[(w)1]} # extract first word (pid)
                }
                _wanted nonk-pids expl 'non-kernel process ID' \
                    compadd -o nosort -ld lines -a pids
            }
            ;;
        (probe)
            _arguments -C \
                $arch_compstr \
                $color_compstr \
                $output_compstr \
                '*:arguments:_files'
            ;;
        (emu)
            _arguments -C \
                $arch_compstr \
                $color_compstr \
                '(-q --quiet)'{-q,--quiet}'[Print return value only]' \
                '2:BPF:_files' \
                ':NR:->syscall_nr' \
                ':ARGV0:->argv0' \
                ':ARGV1:->argv1' \
                ':ARGV2:->argv2' \
                ':ARGV3:->argv3' \
                ':ARGV4:->argv4' \
                ':ARGV5:->argv5' \
                ':IP:->ip' \
                '*: :'
            case $state {
                (syscall_nr)
                    _message -r 'Hint: syscall_nr or syscall_name (to name but a few)'
                    _values SYSCALL $syscalls
                    ;;
                (argv0)
                    _message -r 'Hint: u64 for argv[0]'
                    ;;
                (argv1)
                    _message -r 'Hint: u64 for argv[1]'
                    ;;
                (argv2)
                    _message -r 'Hint: u64 for argv[2]'
                    ;;
                (argv3)
                    _message -r 'Hint: u64 for argv[3]'
                    ;;
                (argv4)
                    _message -r 'Hint: u64 for argv[4]'
                    ;;
                (argv5)
                    _message -r 'Hint: u64 for argv[5]'
                    ;;
                (ip)
                    _message -r 'Hint: u64 for instruction pointer'
                    ;;
                }
            ;;
        (*)
            # help and version has no completion available
            return
            ;;
        }
}
