#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Copyright (c) 2026 Micron Technology, Inc.
"""Generate shell completion scripts for nvme-cli.

Input is the JSON metadata emitted by `nvme utils dump-command-metadata`, read
from a file or stdin. The command/option model is captured at runtime from the
compiled binary, so this script does no C parsing -- it only formats.

Usage:
    nvme utils dump-command-metadata | generate-completions.py \\
        --bash bash-nvme-completion.sh --zsh _nvme
    generate-completions.py model.json --bash -        # '-' writes to stdout
"""

import argparse
import json
import sys


# Metadata schema version this generator understands; reject anything else.
SUPPORTED_SCHEMA_VERSION = 1


# ---------------------------------------------------------------------------
# Model access helpers
# ---------------------------------------------------------------------------

def opt_suffix(opt):
    """'=' if the option takes an argument, else ''."""
    return "=" if opt.get("argument", "none") != "none" else ""


def opt_word(opt):
    """The completion word for an option: '--name=' or '--name'."""
    return "--" + opt["long"] + opt_suffix(opt)


def flatten(s):
    """Collapse whitespace so a multi-line description stays on one line."""
    return " ".join((s or "").split())


def cmd_names(cmd):
    """Every name a command answers to: its primary name plus any alias.

    Names are whitespace-stripped.
    """
    names = [cmd["name"]]
    if cmd.get("alias"):
        names.append(cmd["alias"])
    return [n for n in (s.strip() for s in names) if n]


# Completion values for options the metadata carries none for. Keyed by long
# option name, applied to every command with it, so only list values that are
# uniform across all commands.
# NOTE: This is a stopgap. These should probably be defined in the C option
# definitions (OPT_VALS) so they flow through the metadata for all tooling.
VALUE_HINTS = {
    "output-format-version": ["1", "2"],
    "sel": ["0", "1", "2", "3"],
}

# help and version are matched as literal strings in the C dispatcher, not
# registered in any command table, so dump-command-metadata never reports them.
# Listed here since putting them in the metadata would just hardcode them in C.
BUILTIN_COMMANDS = ["help", "version"]


def command_options(cmd):
    """Yield a command's options (global and local alike).

    Hidden options are always skipped: they are accepted on the command line
    but suppressed from --help, so they should not be offered as completions.

    Options that have no enforced value set but do have a completion hint (see
    VALUE_HINTS) are yielded with that hint filled into "values".
    """
    for opt in cmd.get("options", []):
        if opt.get("hidden"):
            continue
        if "values" not in opt and opt["long"] in VALUE_HINTS:
            opt = dict(opt, values=VALUE_HINTS[opt["long"]])
        yield opt


def plugins(model):
    return model.get("plugins", [])


def builtin_commands(model):
    return model.get("commands", [])


# ---------------------------------------------------------------------------
# bash
# ---------------------------------------------------------------------------

BASH_HEADER = '''\
# SPDX-License-Identifier: GPL-2.0-or-later
#
# bash tab completion for the nvme command line utility
#
# This file is generated by generate-completions.py. Do not edit.

# COMPREPLY=( $(compgen ...) ) is the standard bash-completion idiom; the word
# splitting SC2207 warns about is intentional here.
# shellcheck disable=SC2207

# Helper function to detect if we're completing an option's value.
# Uses: $cur, $prev, $words, $cword (from _init_completion)
# Sets: $opt (the option name), $val (partial value), $completing_value (0 or 1)
# Known limitation: a value beginning with '-' (e.g. '--value -5') is
# indistinguishable from an option name here and is treated as one, so its value
# is not completed. Rare for nvme options and not worth the ambiguity to chase.
_nvme_detect_value_completion() {
\tcompleting_value=0
\topt=""
\tval=""

\tif [[ $cur == -*=* ]]; then
\t\topt="${cur%%=*}"
\t\tval="${cur#*=}"
\t\tcompleting_value=1
\telif [[ $cur == "=" ]] && [[ $prev == -* ]]; then
\t\topt="$prev"
\t\tcompleting_value=1
\telif [[ $cur != -* ]] && [[ $cur != "" ]] && [[ $prev == "=" ]] && [[ ${words[$cword-2]} == -* ]]; then
\t\topt="${words[$cword-2]}"
\t\tval="$cur"
\t\tcompleting_value=1
\telif [[ $cur != -* ]] && [[ $cur != "" ]] && [[ $prev == --* ]]; then
\t\topt="$prev"
\t\tval="$cur"
\t\tcompleting_value=1
\telif [[ $cur == "" ]] && [[ $prev == --* ]]; then
\t\topt="$prev"
\t\tcompleting_value=1
\telif [[ $cur == "" ]] && [[ $prev == "=" ]]; then
\t\topt="${words[$cword-2]}"
\t\tcompleting_value=1
\telif [[ $cur == "" ]] && [[ $prev == -? ]]; then
\t\topt="$prev"
\t\tcompleting_value=1
\telif [[ $cur != -* ]] && [[ $cur != "" ]] && [[ $prev == -? ]]; then
\t\topt="$prev"
\t\tval="$cur"
\t\tcompleting_value=1
\tfi

\t# When the option word was captured unsplit (e.g. '--output-format='), drop
\t# the trailing '=' so $opt matches the option names below.
\topt="${opt%=}"
}

# Complete an option's value from a known set. Called once per command with
# alternating NAMES VALUES arguments: NAMES is the option's spellings
# ('--output-format -o'), VALUES its value list. When completing the value of a
# matching option, appends it to $vals. No-op unless completing a value.
_nvme_opt_vals () {
\t[[ $completing_value -eq 1 ]] || return 0
\tlocal name
\twhile (( $# >= 2 )); do
\t\t# shellcheck disable=SC2086  # $1 is a space-separated spelling list
\t\tfor name in $1; do
\t\t\t[[ $opt == "$name" ]] && { vals+=" $2"; return 0; }
\t\tdone
\t\tshift 2
\tdone
}

# Mark that an option's value is a filename/path. Called once per command with
# each argument an option's spellings ('--firmware -f'). When completing the
# value of a matching option, sets $wantfiles so 'complete -o default' offers
# files. No-op unless completing a value.
_nvme_opt_files () {
\t[[ $completing_value -eq 1 ]] || return 0
\tlocal names name
\tfor names; do
\t\t# shellcheck disable=SC2086  # $names is a space-separated spelling list
\t\tfor name in $names; do
\t\t\t[[ $opt == "$name" ]] && { wantfiles=1; return 0; }
\t\tdone
\tdone
}
'''

BASH_FUNC_PREAMBLE = '''\

{func} () {{
\tlocal opts=""
\tlocal valopts=""
\tlocal vals=""
\tlocal opt=""
\tlocal val=""

\topts+=" "
\tvals+=" "

\tlocal completing_value=0
\tlocal wantfiles=0
\t_nvme_detect_value_completion
'''

# Shared tail for every command-group function. Each function fills in the
# per-command $opts/$valopts/$vals via its own case statement, then calls this
# to finish: append -h/--help, offer the device argument when appropriate, and
# emit COMPREPLY. Reads the caller's locals (opts, valopts, vals, val, opt,
# completing_value, wantfiles, words, cur, COMPREPLY) via bash dynamic scoping,
# the same mechanism _nvme_detect_value_completion relies on. $1 is the command
# name; $2 is the position at which the device argument is expected (2 for
# builtins, 3 for plugin sub-commands).
#
# A word counts as a positional argument only if it is not an option, not a bare
# '=', and not the value of a preceding value-taking option -- in either the
# '--opt val' or split '--opt = val' form -- so e.g. '--output-file /dev/nvme0'
# does not count the path as the device.
BASH_FINISH_FUNC = '''\

_nvme_finish_completion () {
\topts+=" -h --help"

\tif [[ $completing_value -eq 1 ]] && [[ " $valopts " != *" $opt "* ]]; then
\t\tcompleting_value=0
\tfi

\tif [[ $completing_value -eq 0 ]]; then
\t\tlocal nonopt_args=0 has_device=0 i prevopt
\t\tfor (( i=0; i < ${#words[@]}-1; i++ )); do
\t\t\t[[ ${words[i]} == -* || ${words[i]} == "=" ]] && continue
\t\t\t# Skip words[i] if it is the value of a preceding value-taking option.
\t\t\t# That option is the previous word ('--opt val'), or the word before a
\t\t\t# split '=' ('--opt = val').
\t\t\tif [[ $i -gt 0 ]]; then
\t\t\t\tprevopt="${words[i-1]}"
\t\t\t\t[[ $i -ge 2 && $prevopt == "=" ]] && prevopt="${words[i-2]}"
\t\t\t\t# Strip a trailing '=' so an unsplit '--opt=' token still matches
\t\t\t\t# $valopts (whose entries carry no '=').
\t\t\t\t[[ " $valopts " == *" ${prevopt%=} "* ]] && continue
\t\t\tfi
\t\t\t(( nonopt_args += 1 ))
\t\t\t[[ ${words[i]} == /dev/nvme* ]] && has_device=1
\t\tdone
\t\tif [[ $nonopt_args -ge $2 ]] && [[ $has_device -eq 0 ]] && \\
\t\t   [[ "$1" != "help" ]] && [[ "$1" != "version" ]]; then
\t\t\topts="/dev/nvme* $opts"
\t\tfi
\tfi

\tif [[ $completing_value -eq 1 ]]; then
\t\tif [[ $vals != " " ]]; then
\t\t\tCOMPREPLY+=( $( compgen -W "$vals" -- "$val" ) )
\t\telif [[ $wantfiles -eq 1 ]]; then
\t\t\t: # value is a filename; let 'complete -o default' offer files
\t\telse
\t\t\t# No candidate values (e.g. a NUM or free-form string); offer
\t\t\t# nothing and suppress the filename fallback.
\t\t\tcompopt +o default
\t\tfi
\telse
\t\tCOMPREPLY+=( $( compgen -W "$opts" -- "$cur" ) )
\t\t# Suppress the trailing space only when the sole completion ends in '=',
\t\t# so the user can type the value immediately. With more than one candidate
\t\t# bash appends nothing, so nospace must not fire.
\t\t[[ ${#COMPREPLY[@]} -eq 1 && ${COMPREPLY[0]} == *= ]] && compopt -o nospace
\tfi

\treturn 0
}
'''

# Tail of each command-group function: hand off to the shared finisher. $1 is
# the command name (already this function's first parameter); {device_argpos} is
# baked in per group.
BASH_FUNC_EPILOGUE = '''\
\t_nvme_finish_completion "$1" {device_argpos}
}}
'''


def wrap_words(words, indent, width=72):
    """Join words into indented lines, wrapping near `width` columns."""
    lines = []
    line = indent
    for w in words:
        if line != indent and len(line) + 1 + len(w) > width:
            lines.append(line)
            line = indent
        line += (" " if line != indent else "") + w
    if line != indent:
        lines.append(line)
    return "\n".join(lines)


def bash_option_words(opts):
    """Build the ' --name= -s ...' word list for a sequence of options."""
    out = ""
    for o in opts:
        out += " " + opt_word(o)
        if o.get("short"):
            out += " -" + o["short"]
    return out


def bash_valopt_words(opts):
    """Names (long and short, no '=') of the options that take a value.

    The value-completion path keys off this: an option word not listed here is a
    flag, so completing after it must fall back to normal option completion
    rather than sit in value mode offering nothing.
    """
    out = ""
    for o in opts:
        if opt_suffix(o):
            out += " --" + o["long"]
            if o.get("short"):
                out += " -" + o["short"]
    return out


def opt_is_file(opt):
    """True if the option's value is a filename/path we can file-complete."""
    return opt.get("metavar") in ("FILE", "DIRECTORY")


def opt_spellings(opt):
    """An option's spellings as a space-separated string: '--long' or '--long -s'."""
    sel = "--" + opt["long"]
    if opt.get("short"):
        sel += " -" + opt["short"]
    return sel


def bash_option_body(opts):
    """The 'opts+=', 'valopts+=' and value-completion lines for a command's
    option case, indented three tabs. Returns '' when there is nothing to emit
    (no option words and no value completion)."""
    words = bash_option_words(opts)
    val_opts = [o for o in opts if o.get("values")]
    file_opts = [o for o in opts if not o.get("values") and opt_is_file(o)]
    if not words and not val_opts and not file_opts:
        return ""
    out = f'\t\t\topts+="{words}"\n'
    valwords = bash_valopt_words(opts)
    if valwords:
        out += f'\t\t\tvalopts+="{valwords}"\n'
    if val_opts or file_opts:
        out += "\n"
    if val_opts:
        pairs = [f'"{opt_spellings(o)}" "{" ".join(o["values"])}"' for o in val_opts]
        if len(pairs) == 1:
            out += f'\t\t\t_nvme_opt_vals {pairs[0]}\n'
        else:
            # Continuation lines align under the first argument: 3 tabs of body
            # indent + len("_nvme_opt_vals ") spaces.
            cont = "\t\t\t" + " " * len("_nvme_opt_vals ")
            joined = f" \\\n{cont}".join(pairs)
            out += f'\t\t\t_nvme_opt_vals {joined}\n'
    if file_opts:
        args = " ".join(f'"{opt_spellings(o)}"' for o in file_opts)
        out += f'\t\t\t_nvme_opt_files {args}\n'
    return out


def emit_bash_command_group(out, commands, func, device_argpos):
    out.write(BASH_FUNC_PREAMBLE.format(func=func))

    # One case per command, emitting exactly the options that command accepts
    # (global and local together, straight from the metadata). This deliberately
    # repeats the global options in every command rather than factoring them into
    # a shared '*)' clause: global options are NOT uniform across commands (some
    # take none), so a shared clause would offer options a command rejects. Each
    # command carrying its own list keeps the generator correct by construction
    # when the C command definitions change. version/help accept nothing and so
    # get no case; -h/--help is still appended for every command in the epilogue.
    clauses = []
    for c in commands:
        clause_body = bash_option_body(list(command_options(c)))
        if not clause_body:
            continue
        label = "|".join(f'"{n}"' for n in cmd_names(c))
        clauses.append(f'\t\t{label})\n{clause_body}\t\t\t;;\n')
    if clauses:
        out.write('\tcase "$1" in\n')
        out.write("\n".join(clauses))
        out.write('\tesac\n')

    out.write(BASH_FUNC_EPILOGUE.format(device_argpos=device_argpos))


def bash_func_name(plugin):
    return "plugin_%s_opts" % plugin.replace("-", "_")


def generate_bash(model, out):
    out.write(BASH_HEADER)
    out.write(BASH_FINISH_FUNC)

    emit_bash_command_group(out, builtin_commands(model), "nvme_list_opts", 2)
    for p in plugins(model):
        emit_bash_command_group(out, p["commands"], bash_func_name(p["name"]), 3)

    # Dispatcher.
    out.write("\n_nvme_subcmds () {\n"
              "\tlocal cur prev words cword\n"
              "\t_init_completion || return\n"
              "\n"
              "\ttypeset -Ar _plugin_subcmds=(\n")
    for p in plugins(model):
        names = " ".join(n for c in p["commands"] for n in cmd_names(c))
        out.write(f'\t\t[{p["name"]}]="{names}"\n')
    out.write("\t)\n\n\ttypeset -Ar _plugin_funcs=(\n")
    for p in plugins(model):
        out.write(f'\t\t[{p["name"]}]="{bash_func_name(p["name"])}"\n')
    out.write("\t)\n\n\tlocal -a _cmds=(\n")
    cmds = ([n for c in builtin_commands(model) for n in cmd_names(c)] +
            [p["name"] for p in plugins(model)] +
            BUILTIN_COMMANDS)
    out.write(wrap_words(cmds, "\t\t"))
    out.write("\n\t)\n\n"
              "\tlocal func subcmd\n"
              "\tif [[ ${#words[*]} -lt 3 ]]; then\n"
              "\t\tCOMPREPLY+=( $(compgen -W \"${_cmds[*]}\" -- \"$cur\" ) )\n"
              "\telse\n"
              "\t\tfor subcmd in \"${!_plugin_subcmds[@]}\"; do\n"
              "\t\t\tif [[ ${words[1]} == \"$subcmd\" ]]; then\n"
              "\t\t\t\tif [[ ${#words[*]} -lt 4 ]]; then\n"
              "\t\t\t\t\tCOMPREPLY+=( $(compgen -W \"${_plugin_subcmds[$subcmd]}\" -- \"$cur\" ) )\n"
              "\t\t\t\telse\n"
              "\t\t\t\t\tfunc=${_plugin_funcs[$subcmd]}\n"
              "\t\t\t\t\t$func \"${words[2]}\" \"$prev\"\n"
              "\t\t\t\tfi\n"
              "\t\t\t\treturn 0\n"
              "\t\t\tfi\n"
              "\t\tdone\n"
              "\n"
              "\t\tnvme_list_opts \"${words[1]}\" \"$prev\"\n"
              "\tfi\n"
              "\n"
              "\treturn 0\n"
              "}\n"
              "\ncomplete -o default -F _nvme_subcmds nvme\n")


# ---------------------------------------------------------------------------
# zsh
# ---------------------------------------------------------------------------

ZSH_HEADER = '''\
#compdef _nvme nvme
# SPDX-License-Identifier: GPL-2.0-or-later
#
# zsh completions for the nvme command-line interface
#
# This file is generated by generate-completions.py. Do not edit.
#
# Note: with verbose completion listings (the default), zsh may report more
# "possibilities" than the number of lines it displays -- e.g. "452
# possibilities" for ~153 commands. This is expected, not duplication: when
# descriptions are shown, _describe lays the matches out in a multi-column grid
# and pads it with empty alignment matches, which count toward the total but
# render as nothing. The count scales with terminal width; the displayed
# entries are correct.
'''

# Helpers emitted inside the _nvme function body (hence the tab indent). Each
# command carries its own options straight from the metadata.
ZSH_HELPERS = '''\
\t# Describe options, giving those ending in '=' no trailing space.
\t#
\t# The tags are 'nvme-noeq-options'/'nvme-eq-options', deliberately NOT the
\t# standard 'options' tag: when a device candidate is also offered (compadd
\t# below), the completion system's handling of the reserved 'options' tag
\t# suppressed the whole group at an empty prefix, so '--help' etc. only appeared
\t# once a '-' was typed. A non-reserved tag sidesteps that.
\t_nvme_describe() {
\t\tlocal desc="$1"
\t\tlocal -a opts=("${(@P)2}")
\t\tlocal -a eq_opts noeq_opts
\t\tlocal opt

\t\tfor opt in "${opts[@]}"; do
\t\t\tif [[ "${opt%%:*}" == *= ]]; then
\t\t\t\teq_opts+=("$opt")
\t\t\telse
\t\t\t\tnoeq_opts+=("$opt")
\t\t\tfi
\t\tdone

\t\t(( ${#noeq_opts} )) && _describe -t nvme-noeq-options "$desc" noeq_opts
\t\t(( ${#eq_opts} )) && _describe -t nvme-eq-options "$desc" eq_opts -S ''
\t}

\t# Complete an option's value from a known set, in both the attached
\t# ('--opt=val', '-o=val') and separate-word ('--opt val', '-o val') forms.
\t# Called once per command with alternating NAMES VALUES arguments: NAMES is a
\t# space-separated spelling list ('--output-format -o'), VALUES its
\t# space-separated value list. Completes the first option whose spelling matches
\t# the word being entered.
\t_nvme_opt_vals() {
\t\tlocal -a names
\t\tlocal name
\t\twhile (( $# >= 2 )); do
\t\t\tnames=(${=1})
\t\t\tfor name in "${names[@]}"; do
\t\t\t\tif [[ ${words[CURRENT]} == ${name}=* ]] || \\
\t\t\t\t   [[ ${words[CURRENT-1]} == "${name}" ]]; then
\t\t\t\t\t_values '' ${=2}
\t\t\t\t\treturn 0
\t\t\t\tfi
\t\t\tdone
\t\t\tshift 2
\t\tdone
\t\treturn 1
\t}

\t# Complete an option's value as a filename, in both the attached
\t# ('--opt=path', '-o=path') and separate-word ('--opt path', '-o path') forms.
\t# Each argument is a space-separated spelling list ('--fw -f'); if the word
\t# being entered is that option's value, offer files. Mirrors the value form of
\t# _nvme_opt_vals.
\t_nvme_opt_files() {
\t\tlocal -a names
\t\tlocal name spec
\t\tfor spec in "$@"; do
\t\t\tnames=(${=spec})
\t\t\tfor name in "${names[@]}"; do
\t\t\t\tif [[ ${words[CURRENT]} == ${name}=* ]]; then
\t\t\t\t\tcompset -P '*='
\t\t\t\t\t_files
\t\t\t\t\treturn 0
\t\t\t\telif [[ ${words[CURRENT-1]} == "${name}" ]]; then
\t\t\t\t\t_files
\t\t\t\t\treturn 0
\t\t\t\tfi
\t\t\tdone
\t\tdone
\t\treturn 1
\t}

\t# True if a device argument is already present, so we do not offer another.
\t# Only words before the one being completed count (the current word may be a
\t# partial '/dev/nvme' the user is still typing).
\t_nvme_has_device() {
\t\tlocal -i i
\t\tfor (( i = 1; i < CURRENT; i++ )); do
\t\t\t[[ ${words[i]} == /dev/nvme* ]] && return 0
\t\tdone
\t\treturn 1
\t}
'''


def zsh_escape(s):
    """Escape a string for inclusion inside a zsh '...' single-quoted string."""
    return (s or "").replace("'", "'\\''")


def zsh_var(plugin, cmd):
    """A unique zsh array-variable name for a command's option list."""
    name = "_"
    if plugin:
        name += plugin.replace("-", "_") + "_"
    name += cmd.replace("-", "_")
    return name


def emit_zsh_value_completions(out, tab, cmd):
    """Emit a single _nvme_opt_vals call listing the command's value-set options
    as alternating spelling/value-list pairs."""
    pairs = [f'"{opt_spellings(o)}" "{" ".join(o["values"])}"'
             for o in command_options(cmd) if o.get("values")]
    if not pairs:
        return
    if len(pairs) == 1:
        out.write(f'{tab}_nvme_opt_vals {pairs[0]} && return\n')
    else:
        # Continuation lines align under the first argument.
        cont = tab + " " * len("_nvme_opt_vals ")
        joined = f" \\\n{cont}".join(pairs)
        out.write(f'{tab}_nvme_opt_vals {joined} && return\n')


def emit_zsh_file_completions(out, tab, cmd):
    """Emit a _nvme_opt_files call for the command's file/directory-valued
    options (those without an enumerated value set)."""
    args = [f'"{opt_spellings(o)}"' for o in command_options(cmd)
            if not o.get("values") and opt_is_file(o)]
    if not args:
        return
    if len(args) == 1:
        out.write(f'{tab}_nvme_opt_files {args[0]} && return\n')
    else:
        cont = tab + " " * len("_nvme_opt_files ")
        joined = f" \\\n{cont}".join(args)
        out.write(f'{tab}_nvme_opt_files {joined} && return\n')


def emit_zsh_command_clause(out, tab, path, cmd):
    """Emit the case clause for one command: its value completions, option
    descriptions, and a device-argument candidate. Aliases share the clause."""
    names = cmd_names(cmd)
    # The primary (whitespace-stripped) name drives the variable and label; a
    # stray space in the C command tables would otherwise leak into both.
    name = names[0]
    # The pattern sits at 'tab'; its body is indented one level deeper.
    body = tab + "\t"
    out.write(f'{tab}({"|".join(names)})\n')

    opts = list(command_options(cmd))
    if not opts:
        # No options to offer, but still complete a device argument.
        out.write(f"{body}_nvme_has_device || compadd -- /dev/nvme*(N)\n")
        out.write(f"{body};;\n\n")
        return

    emit_zsh_value_completions(out, body, cmd)
    emit_zsh_file_completions(out, body, cmd)

    var = zsh_var(path, name)
    out.write(f"{body}local {var}\n{body}{var}=(\n")
    for o in opts:
        desc = zsh_escape(flatten(o.get("description")))
        out.write(f"{body}\t--{o['long']}{opt_suffix(o)}':{desc}'\n")
        if o.get("short"):
            out.write(f"{body}\t-{o['short']}':alias for --{o['long']}'\n")
    # -h/--help is accepted by every command but is not in the metadata option
    # list, so add it explicitly.
    out.write(f"{body}\t--help':show help'\n")
    out.write(f"{body}\t-h':alias for --help'\n")
    out.write(f"{body})\n")
    out.write(f"{body}_arguments '*:: :->subcmds'\n")
    plabel = f"nvme {path + ' ' if path else ''}{name} options"
    out.write(f'{body}_nvme_describe "{plabel}" {var}\n')
    out.write(f"{body}_nvme_has_device || compadd -- /dev/nvme*(N)\n")
    out.write(f"{body};;\n\n")


def generate_zsh(model, out):
    out.write(ZSH_HEADER)
    out.write("\n_nvme () {\n\tlocal -a _cmds\n\t_cmds=(\n")
    for c in builtin_commands(model):
        for n in cmd_names(c):
            out.write(f"\t\t'{n}:{zsh_escape(flatten(c.get('description')))}'\n")
    for p in plugins(model):
        out.write(f"\t\t'{p['name']}:{zsh_escape(flatten(p.get('description')))}'\n")
    for name in BUILTIN_COMMANDS:
        out.write(f"\t\t'{name}'\n")
    out.write("\t)\n\n")

    out.write(ZSH_HELPERS)

    # After the top-level _arguments strips 'nvme', CURRENT counts words within
    # the command: 1 = completing the command name, 2 = a builtin's first
    # option OR a plugin's sub-command name, >=3 = a plugin sub-command's
    # options. Builtin option completion and plugin sub-command listing both
    # happen at CURRENT==2, so the dispatch keys on ${words[1]} without a lower
    # CURRENT guard; each plugin clause splits sub-command listing (CURRENT==2)
    # from sub-command option completion itself.
    out.write("\n\t_arguments '*:: :->subcmds' && return 0\n\n"
              "\tif (( CURRENT == 1 )); then\n"
              "\t\t_describe -t commands \"nvme subcommands\" _cmds\n"
              "\t\treturn\n"
              "\tfi\n\n"
              "\tcase ${words[1]} in\n")

    for p in plugins(model):
        out.write(f"\t\t({p['name']})\n"
                  f"\t\t\tif (( CURRENT == 2 )); then\n"
                  f"\t\t\t\tlocal -a _sub\n\t\t\t\t_sub=(\n")
        for c in p["commands"]:
            desc = zsh_escape(flatten(c.get("description")))
            for n in cmd_names(c):
                out.write(f"\t\t\t\t\t'{n}:{desc}'\n")
        out.write(f"\t\t\t\t)\n"
                  f'\t\t\t\t_describe -t commands "nvme {p["name"]} commands" _sub\n'
                  f"\t\t\t\treturn\n\t\t\tfi\n\n"
                  f"\t\t\tcase ${{words[2]}} in\n")
        for c in p["commands"]:
            emit_zsh_command_clause(out, "\t\t\t\t", p["name"], c)
        out.write("\t\t\t\t(*)\n\t\t\t\t\t_files\n\t\t\t\t\t;;\n")
        out.write("\t\t\tesac\n\t\t\t;;\n\n")

    for c in builtin_commands(model):
        emit_zsh_command_clause(out, "\t\t", None, c)

    # help and version have no metadata command clause. 'help CMD' prints CMD's
    # help, so complete a command name at that one argument (CURRENT==2); past it
    # there is nothing to complete. 'version' takes no argument. Both return
    # rather than fall through to _files.
    out.write('\t\t(help)\n'
              '\t\t\t(( CURRENT == 2 )) && \\\n'
              '\t\t\t\t_describe -t commands "nvme subcommands" _cmds\n'
              '\t\t\treturn\n\t\t\t;;\n\n'
              '\t\t(version)\n'
              '\t\t\treturn\n\t\t\t;;\n\n')

    out.write("\t\t(*)\n\t\t\t_files\n\t\t\t;;\n"
              "\tesac\n}\n")


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

GENERATORS = {"bash": generate_bash, "zsh": generate_zsh}


def open_out(path):
    return sys.stdout if path == "-" else open(path, "w")


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input", nargs="?", default="-",
                    help="JSON model file, or '-' for stdin (default)")
    ap.add_argument("--bash", metavar="FILE",
                    help="write bash completion to FILE ('-' for stdout)")
    ap.add_argument("--zsh", metavar="FILE",
                    help="write zsh completion to FILE ('-' for stdout)")
    args = ap.parse_args()

    requested = [(s, getattr(args, s)) for s in GENERATORS if getattr(args, s)]
    if not requested:
        ap.error("specify at least one of --bash/--zsh")

    # Each shell's script is distinct, so several sharing stdout would just
    # concatenate into one unusable stream. Require a named file per shell when
    # more than one is requested.
    if len(requested) > 1 and any(path == "-" for _, path in requested):
        ap.error("with more than one shell, give each a FILE (not '-'/stdout)")

    if args.input == "-":
        model = json.load(sys.stdin)
    else:
        with open(args.input) as src:
            model = json.load(src)

    version = model.get("schema_version")
    if version != SUPPORTED_SCHEMA_VERSION:
        ap.error(f"unsupported schema_version {version!r} "
                 f"(expected {SUPPORTED_SCHEMA_VERSION})")

    for shell, path in requested:
        out = open_out(path)
        try:
            GENERATORS[shell](model, out)
        finally:
            if out is not sys.stdout:
                out.close()


if __name__ == "__main__":
    main()
