/// <reference types="tree-sitter-cli/dsl" />
// @ts-check

const PREC = {
  RANGE: 1,
  FALLBACK: 2,
  OR: 3,
  AND: 4,
  COMPARE: 5,
  BITWISE: 6,
  SHIFT: 7,
  SUM: 8,
  PRODUCT: 9,
  PREFIX: 10,
  POSTFIX: 11,
};

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

  word: $ => $.identifier,

  extras: $ => [/[ \t\r]/, $.comment],

  conflicts: $ => [
    [$.expression, $.qualified_identifier],
    [$.constant_declaration, $.variable_declaration, $.expression],
    [$.constant_declaration, $.variable_declaration, $.expression, $.qualified_identifier],
    [$.function_declaration],
    [$.if_statement],
    [$.enum_literal],
    [$.array_type],
    [$.array_type, $.expression],
    [$.array_type, $.array_literal],
    [$.expression_statement, $.parenthesized_expression],
    [$.block, $.tuple_literal, $.anonymous_record_literal],
    [$.field_initializer, $.expression],
    [$.capture_list, $.expression],
    [$.match_capture, $.expression],
  ],

  rules: {
    source_file: $ => repeat(choice($._newline, $._top_level_declaration)),

    _top_level_declaration: $ => choice(
      $.import_declaration,
      $.test_import_declaration,
      $.function_declaration,
      $.test_declaration,
      $.type_declaration,
      $.global_constant_declaration,
      $.global_variable_declaration,
    ),

    import_declaration: $ => seq(
      optional(seq(field('alias', $.identifier), '::', repeat($._newline))),
      'import',
      repeat($._newline),
      field('path', $.string),
    ),

    test_import_declaration: $ => seq(
      'test',
      'import',
      repeat($._newline),
      field('path', $.string),
    ),

    test_declaration: $ => seq(
      field('name', $.identifier),
      'test',
      repeat($._newline),
      field('body', $.block),
    ),

    function_declaration: $ => seq(
      optional('hide'),
      field('name', $.identifier),
      field('kind', choice('func', 'c_func')),
      field('parameters', $.parameter_list),
      repeat($._newline),
      field('result', $.type),
      optional(seq('!', field('error', $._error_type))),
      optional(choice(
        field('body', $.block),
        seq(repeat1($._newline), field('body', $.block)),
      )),
    ),

    type_declaration: $ => prec(1, seq(
      optional('hide'),
      field('name', $.identifier),
      '::',
      repeat($._newline),
      field('value', choice(
        $.struct_type,
        $.c_struct_type,
        $.union_type,
        $.enum_type,
        $.opaque_type,
        $.distinct_type,
        $.alias_type,
      )),
    )),

    global_constant_declaration: $ => seq(
      optional('hide'),
      field('name', $.identifier),
      optional(field('type', $.type)),
      '::',
      repeat($._newline),
      field('value', $._value),
    ),

    global_variable_declaration: $ => seq(
      optional('hide'),
      field('name', $.identifier),
      optional(field('type', $.type)),
      '=',
      repeat($._newline),
      field('value', $._value),
    ),

    struct_type: $ => seq('struct', repeat($._newline), $.record_body),
    c_struct_type: $ => seq('c_struct', repeat($._newline), $.record_body),
    opaque_type: _ => 'opaque',
    distinct_type: $ => seq('distinct', field('type', $.type)),
    alias_type: $ => seq('alias', field('type', $.type)),

    union_type: $ => seq(
      'union',
      optional(seq('(', choice('enum', $.type), ')')),
      repeat($._newline),
      $.record_body,
    ),

    enum_type: $ => seq(
      'enum',
      optional(seq('(', field('backing', $.type), ')')),
      repeat($._newline),
      $.enum_body,
    ),

    record_body: $ => seq(
      '{',
      repeat(choice($._newline, seq($.record_field, optional(',')))),
      '}',
    ),

    record_field: $ => choice(
      seq(
        field('name', $.identifier),
        field('type', choice($.type, $.struct_type)),
        optional(seq('=', repeat($._newline), field('default', $.expression))),
      ),
      field('type', choice($.type, $.struct_type)),
    ),

    enum_body: $ => seq(
      '{',
      repeat(choice($._newline, seq($.enum_member, optional(',')))),
      '}',
    ),

    enum_member: $ => seq(
      field('name', $.identifier),
      optional(seq('=', optional('-'), field('value', $.integer))),
    ),

    parameter_list: $ => seq(
      '(',
      commaSep($, choice($.parameter, '...')),
      ')',
    ),

    parameter: $ => seq(
      optional('$'),
      field('name', choice($.identifier, $.sink)),
      repeat(seq(',', repeat($._newline), optional('$'), field('name', choice($.identifier, $.sink)))),
      field('type', $.type),
    ),

    type: $ => prec.left(seq(
      $._type_atom,
      repeat(seq('|', $._type_atom)),
    )),

    _type_atom: $ => choice(
      $.builtin_type,
      $.intrinsic_type,
      $.named_type,
      $.optional_type,
      $.pointer_type,
      $.array_type,
      $.function_type,
      $.parenthesized_type,
    ),

    parenthesized_type: $ => seq('(', repeat($._newline), $.type, repeat($._newline), ')'),

    named_type: $ => prec.right(seq($.qualified_identifier, optional($.argument_list))),

    intrinsic_type: $ => prec(PREC.POSTFIX, seq(
      field('function', alias('struct_type!', $.identifier)),
      field('arguments', $.argument_list),
    )),

    optional_type: $ => seq('?', $.type),

    pointer_type: $ => seq(
      choice('@', '*'),
      optional('mut'),
      $.type,
    ),

    array_type: $ => seq(
      '[',
      repeat($._newline),
      optional(choice(
        seq('*', ';', field('sentinel', $._type_constant)),
        seq(';', field('sentinel', $._type_constant)),
        seq(
          field('length', choice($.sink, $.expression)),
          optional(seq(';', field('sentinel', $._type_constant))),
        ),
      )),
      repeat($._newline),
      ']',
      optional('mut'),
      field('element', $.type),
    ),

    function_type: $ => prec.right(seq(
      choice('func', 'c_func'),
      $.parameter_list,
      repeat($._newline),
      field('result', $.type),
      optional(seq('!', field('error', $._error_type))),
    )),

    _error_type: $ => choice($.type, $.enum_type, $.union_type),
    _type_constant: $ => seq(optional('-'), choice($.integer, $.character)),

    builtin_type: _ => choice(
      'void', 'noreturn', 'type', 'anyopaque', 'bool', 'int', 'uint', 'float', 'range',
      'i8', 'i16', 'i32', 'i64', 'u8', 'u16', 'u32', 'u64',
      'isize', 'usize', 'f32', 'f64',
      'c_char', 'c_schar', 'c_uchar', 'c_short', 'c_ushort',
      'c_int', 'c_uint', 'c_long', 'c_ulong', 'c_longlong',
      'c_ulonglong', 'c_float', 'c_double', 'c_longdouble',
    ),

    block: $ => seq(
      '{',
      repeat(choice($._newline, $.statement)),
      '}',
    ),

    statement: $ => choice(
      $.constant_declaration,
      $.variable_declaration,
      $.assignment_statement,
      $.return_statement,
      $.yield_statement,
      $.if_statement,
      $.while_statement,
      $.for_statement,
      $.match_statement,
      $.break_statement,
      $.continue_statement,
      $.defer_statement,
      $.labeled_block,
      $.block,
      $.expression_statement,
    ),

    constant_declaration: $ => seq(
      field('name', choice($.identifier, $.sink)),
      optional(field('type', $.type)),
      '::',
      repeat($._newline),
      field('value', $._value),
    ),

    variable_declaration: $ => seq(
      field('name', choice($.identifier, $.sink)),
      field('type', $.type),
      '=',
      repeat($._newline),
      field('value', $._value),
    ),

    assignment_statement: $ => seq(
      field('left', $.expression),
      field('operator', choice('=', '+=', '-=', '*=', '/=', '&=', '|=', 'xor=', '<<=', '>>=', '<<|=')),
      repeat($._newline),
      field('right', $._value),
    ),

    expression_statement: $ => $.expression,

    return_statement: $ => prec.right(seq(
      'return',
      optional(field('value', $._value)),
    )),

    yield_statement: $ => seq(
      'yield',
      optional(seq(':', field('label', $.identifier))),
      field('value', $._value),
    ),

    break_statement: $ => seq('break', optional(seq(':', field('label', $.identifier)))),
    continue_statement: $ => seq('continue', optional(seq(':', field('label', $.identifier)))),

    defer_statement: $ => choice(
      seq('defer', repeat($._newline), field('body', $.statement)),
      seq(
        'errdefer',
        repeat($._newline),
        optional(seq(field('capture', $.error_capture), repeat($._newline))),
        field('body', $.statement),
      ),
    ),

    error_capture: $ => seq('|', choice($.identifier, $.sink), '|'),

    labeled_block: $ => seq(
      field('label', $.identifier),
      ':',
      repeat($._newline),
      $.block,
    ),

    if_statement: $ => seq(
      'if',
      repeat($._newline),
      field('condition', $.expression),
      optional($.capture_list),
      repeat($._newline),
      field('consequence', $._branch_body),
      optional(seq(
        repeat($._newline),
        'else',
        repeat($._newline),
        field('alternative', $._branch_body),
      )),
    ),

    capture_list: $ => seq(
      '|',
      commaSep1($, choice($.identifier, $.sink)),
      optional(seq(':', field('guard', $.expression))),
      '|',
    ),

    while_statement: $ => seq(
      'while',
      repeat($._newline),
      field('condition', $.expression),
      optional(seq(':', repeat($._newline), field('update', choice(
        $.assignment_statement,
        $.expression_statement,
        seq('(', repeat($._newline), choice($.assignment_statement, $.expression_statement), repeat($._newline), ')'),
      )))),
      repeat($._newline),
      optional(seq(field('label', $.identifier), ':', repeat($._newline))),
      field('body', $.block),
    ),

    for_statement: $ => seq(
      optional('expand'),
      'for',
      repeat($._newline),
      field('iterable', $.expression),
      repeat($._newline),
      $._for_capture_bar,
      optional('@'),
      field('item', $.identifier),
      optional(seq(',', field('index', $.identifier))),
      $._for_capture_bar,
      repeat($._newline),
      optional(seq(field('label', $.identifier), ':', repeat($._newline))),
      field('body', $.block),
    ),

    _for_capture_bar: _ => token(prec(1, '|')),

    match_statement: $ => seq(
      'match',
      repeat($._newline),
      field('subject', $.expression),
      repeat($._newline),
      '{',
      repeat(choice($._newline, $.match_arm)),
      '}',
    ),

    match_arm: $ => choice(seq(
      field('pattern', choice('else', commaSep1($, $.expression))),
      optional($.match_capture),
      ':',
      repeat($._newline),
      field('body', $._branch_body),
    ), seq(
      'expand',
      $.expand_match_capture,
      ':',
      repeat($._newline),
      field('body', $._branch_body),
    )),

    match_capture: $ => seq('|', optional('@'), field('name', choice($.identifier, $.sink)), '|'),

    expand_match_capture: $ => seq(
      '|',
      optional('@'),
      field('value', choice($.identifier, $.sink)),
      optional(seq(',', field('tag', choice($.identifier, $.sink)))),
      '|',
    ),

    _branch_body: $ => $.statement,

    _value: $ => choice(
      $.labeled_block,
      $.block,
      $.if_statement,
      $.while_statement,
      $.for_statement,
      $.match_statement,
      $.expression,
    ),

    expression: $ => choice(
      $.binary_expression,
      $.catch_expression,
      $.unary_expression,
      $.field_expression,
      $.intrinsic_call_expression,
      $.call_expression,
      $.index_expression,
      $.slice_expression,
      $.postfix_expression,
      $.struct_literal,
      $.anonymous_record_literal,
      $.tuple_literal,
      $.comptime_block,
      $.function_literal,
      $.struct_type,
      $.array_type,
      $.enum_literal,
      $.array_literal,
      $.parenthesized_expression,
      $.identifier,
      $.sink,
      $.builtin_type,
      $.integer,
      $.float,
      $.string,
      $.multiline_string,
      $.character,
      $.boolean,
      $.null,
      $.unreachable,
      $.undefined,
    ),

    binary_expression: $ => choice(
      prec.left(PREC.RANGE, seq(field('left', $.expression), field('operator', choice('..', '..=')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.FALLBACK, seq(field('left', $.expression), field('operator', choice('orelse', 'catch')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.OR, seq(field('left', $.expression), 'or', repeat($._newline), field('right', $.expression))),
      prec.left(PREC.AND, seq(field('left', $.expression), 'and', repeat($._newline), field('right', $.expression))),
      prec.left(PREC.COMPARE, seq(field('left', $.expression), field('operator', choice('==', '!=', '<', '<=', '>', '>=')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.BITWISE, seq(field('left', $.expression), field('operator', choice('&', 'xor', '|')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.SHIFT, seq(field('left', $.expression), field('operator', choice('<<', '>>', '<<|')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.SUM, seq(field('left', $.expression), field('operator', choice('+', '-')), repeat($._newline), field('right', $.expression))),
      prec.left(PREC.PRODUCT, seq(field('left', $.expression), field('operator', choice('*', '/')), repeat($._newline), field('right', $.expression))),
    ),

    catch_expression: $ => prec.left(PREC.FALLBACK, seq(
      field('value', $.expression),
      'catch',
      '|',
      field('name', choice($.identifier, $.sink)),
      '|',
      repeat($._newline),
      field('body', $.block),
    )),

    unary_expression: $ => prec(PREC.PREFIX, seq(
      field('operator', choice('-', '&', '!', '~', '$', 'try')),
      repeat($._newline),
      field('operand', $.expression),
    )),

    field_expression: $ => prec.left(PREC.POSTFIX, seq(
      field('value', $.expression),
      '.',
      field('field', choice($.identifier, $.integer)),
    )),

    intrinsic_call_expression: $ => prec(PREC.POSTFIX, seq(
      field('function', $.identifier),
      field('marker', '!'),
      field('arguments', $.argument_list),
    )),

    call_expression: $ => prec.left(PREC.POSTFIX, seq(
      field('function', $.expression),
      field('arguments', $.argument_list),
    )),

    argument_list: $ => prec(PREC.POSTFIX, seq('(', commaSep($, $.expression), ')')),

    index_expression: $ => prec.left(PREC.POSTFIX, seq(
      field('value', $.expression),
      '[',
      repeat($._newline),
      field('index', $.expression),
      repeat($._newline),
      ']',
    )),

    slice_expression: $ => prec.left(PREC.POSTFIX, seq(
      field('value', $.expression),
      '[',
      repeat($._newline),
      optional(field('start', $.expression)),
      '..',
      optional(field('end', $.expression)),
      repeat($._newline),
      ']',
    )),

    postfix_expression: $ => prec.left(PREC.POSTFIX, seq(
      field('value', $.expression),
      field('operator', choice('?', '^')),
    )),

    struct_literal: $ => prec(PREC.POSTFIX, seq(
      field('type', $.qualified_identifier),
      optional($.argument_list),
      field('fields', $.initializer_list),
    )),

    initializer_list: $ => seq(
      '{',
      choice(
        repeat($._newline),
        seq(
          repeat($._newline),
          choice($.field_initializer, $.expression),
          repeat(seq(repeat($._newline), ',', repeat($._newline), choice($.field_initializer, $.expression))),
          optional(seq(repeat($._newline), ',')),
          repeat($._newline),
        ),
      ),
      '}',
    ),

    field_initializer: $ => seq(
      field('name', $._field_name),
      optional(seq('=', repeat($._newline), field('value', $.expression))),
    ),

    anonymous_record_literal: $ => prec(2, seq(
      '{',
      repeat($._newline),
      $.keyed_field_initializer,
      repeat(seq(repeat($._newline), ',', repeat($._newline), $.keyed_field_initializer)),
      optional(seq(repeat($._newline), ',')),
      repeat($._newline),
      '}',
    )),

    tuple_literal: $ => prec(1, choice(
      seq('{', repeat($._newline), '}'),
      seq(
        '{', repeat($._newline), $.expression,
        repeat(seq(repeat($._newline), ',', repeat($._newline), $.expression)),
        optional(seq(repeat($._newline), ',')),
        repeat($._newline), '}',
      ),
    )),

    enum_literal: $ => seq(
      '.',
      field('name', $.identifier),
      optional($.variant_payload),
    ),

    variant_payload: $ => seq(
      '{',
      choice(
        repeat($._newline),
        seq(
          repeat($._newline),
          choice(
            $.expression,
            seq(
              $.keyed_field_initializer,
              repeat(seq(repeat($._newline), ',', repeat($._newline), $.keyed_field_initializer)),
              optional(seq(repeat($._newline), ',')),
            ),
          ),
          repeat($._newline),
        ),
      ),
      '}',
    ),

    keyed_field_initializer: $ => seq(
      field('name', $._field_name),
      '=',
      repeat($._newline),
      field('value', $.expression),
    ),

    array_literal: $ => seq('[', commaSep($, $.expression), ']'),

    parenthesized_expression: $ => seq(
      '(',
      repeat($._newline),
      $.expression,
      repeat($._newline),
      ')',
    ),

    comptime_block: $ => seq('$', repeat($._newline), $.block),

    function_literal: $ => seq(
      'func',
      $.parameter_list,
      repeat($._newline),
      field('result', $.type),
      optional(seq('!', field('error', $._error_type))),
      repeat($._newline),
      field('body', $.block),
    ),

    qualified_identifier: $ => prec.right(seq(
      field('qualifier', $.identifier),
      optional(seq('.', field('name', $.identifier))),
    )),

    boolean: _ => choice('true', 'false'),
    null: _ => 'null',
    unreachable: _ => 'unreachable',
    undefined: _ => 'undefined',
    sink: _ => '_',

    _field_name: $ => prec(2, choice(
      $.identifier,
      alias(choice(
        'test', 'func', 'c_func', 'struct', 'c_struct', 'opaque', 'union', 'enum',
        'distinct', 'alias', 'import', 'hide', 'return', 'try', 'catch', 'mut',
        'null', 'unreachable', 'undefined', 'orelse', 'and', 'or', 'xor', 'if', 'while', 'for',
        'expand', 'break', 'continue', 'defer', 'errdefer', 'yield', 'match', 'else',
        'true', 'false',
        'void', 'noreturn', 'type', 'anyopaque', 'bool', 'int', 'uint', 'float', 'range',
        'i8', 'i16', 'i32', 'i64', 'u8', 'u16', 'u32', 'u64',
        'isize', 'usize', 'f32', 'f64',
        'c_char', 'c_schar', 'c_uchar', 'c_short', 'c_ushort',
        'c_int', 'c_uint', 'c_long', 'c_ulong', 'c_longlong',
        'c_ulonglong', 'c_float', 'c_double', 'c_longdouble',
      ), $.identifier),
    )),

    identifier: _ => /[A-Za-z_][A-Za-z0-9_]*/,
    integer: _ => /[0-9]+/,
    float: _ => token(prec(1, /[0-9]+\.[0-9]+/)),

    string: $ => seq(
      '"',
      repeat(choice($.string_content, $.escape_sequence)),
      '"',
    ),
    string_content: _ => token.immediate(prec(1, /[^"\\\n]+/)),
    escape_sequence: _ => token.immediate(/\\(?:\\|"|n|r|t|0)/),
    multiline_string: _ => token(/`[^\n]*(?:\n[ \t]*`[^\n]*)*/),
    character: _ => token(/'(?:[^'\\\n]|\\(?:\\|'|n|r|t|0))'/),
    comment: _ => token(seq('#', /[^\n]*/)),
    _newline: _ => /\n/,
  },
});

function commaSep($, rule) {
  return choice(
    repeat($._newline),
    seq(
      repeat($._newline),
      rule,
      repeat(seq(repeat($._newline), ',', repeat($._newline), rule)),
      optional(seq(repeat($._newline), ',')),
      repeat($._newline),
    ),
  );
}

function commaSep1($, rule) {
  return seq(
    rule,
    repeat(seq(repeat($._newline), ',', repeat($._newline), rule)),
  );
}