module.exports = grammar({
  name: 'nix_shell',

  rules: {
    // TODO: allow trailing whitespace on any line
    source_file: $ => seq($.first_line, repeat($.next_line)),

    first_line: $ => seq($._shebang_open, $._opt_ws, optional(seq($._env, $._ws)), $._nix_shell),
    next_line: $ => seq($._newline, choice(
        $.annotation_line,
        repeat(/./),
    )),
    annotation_line: $ => seq($._shebang_open, $._opt_ws, $._nix_shell, optional($.nix_shell_args)),

    // nix acts as if it actually just `exec`'s whichever line contains `#! nix-shell` in the style of `sh`'s `exec`.
    // so arguments can be quoted, invoke subshells, interpolate environment variables, etc.
    // the full scope cannot be covered, but the minimal scope here gets 99% of use cases.
    nix_shell_args: $ => repeat1(seq($._ws, $._nix_shell_arg)),
    _nix_shell_arg: $ => choice(
      seq('-i', $._opt_ws, $.interpreter),
      $.sh_arg,
    ),
    interpreter: $ => $.sh_arg,
    sh_arg: $ => choice(
      $._sh_lit,
      $._sh_quote1,
      $._sh_quote2,
    ),
    _sh_lit: $ => repeat1(choice(
      /[^ \\]/,
      seq('\\', /./),
    )),
    _sh_quote1: $ => seq(
      '\'',
      repeat(choice(
        /[^'\\]/,
        seq('\\', /./),
      )),
      '\'',
    ),
    _sh_quote2: $ => seq(
      '"',
      repeat(choice(
        /[^"\\]/,
        seq('\\', /./),
      )),
      '"',
    ),

    _shebang_open: $ => '#!',
    // TODO: env accepts flags like `-v` or `--unset=NAME` or `VAR=VALUE` before `nix-shell`,
    _env: $ => choice(
      '/usr/bin/env',
      '/bin/env',
      'env',
    ),
    _nix_shell: $ => 'nix-shell',
    // N.B.: this accepts more than it needs to:
    // - shebang parser allows tab characters
    // - #!nix-shell directives do not support tabs
    // - wherever #!nix-shell directives allow whitespace, it's a *single* space -- never multiple
    _ws: $ => /[ \t]+/,
    _opt_ws: $ => /[ \t]*/,
    _newline: $ => '\n',
  }
});