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

  rules: {
    source_file: $ => seq(
      optional($._first_line),
      repeat($._next_line),
    ),
    _first_line: $ => choice(
      $.nix_shell_directive,
      repeat1(/./),
    ),
    _next_line: $ => seq('\n', $._first_line),
    nix_shell_directive: $ => seq(
      /#![ \t]*nix-shell/,
      repeat(seq($._ws, optional($._nix_shell_arg))),
    ),

    // 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_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('\\', /./),
      )),
      '"',
    ),

    // N.B.: this accepts more than it needs to:
    // - #!nix-shell directives do not support tabs
    // - some places where whitespace is accepted (e.g. `#! nix-shell`), only a *single* space is accepted.
    // it's not clear how much of this is intentional, v.s. oversight in matching broader shell parsing,
    // so err on overaccepting in case future nix-shell supports these.
    _ws: $ => /[ \t]+/,
    _opt_ws: $ => /[ \t]*/,
  }
});