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

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

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

  word: $ => $.identifier,

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

  conflicts: $ => [
    [$.expression, $.qualified_identifier],
    [$.constant_declaration, $.variable_declaration, $.expression],
    [$.function_declaration],
    [$.if_statement],
    [$.enum_literal],
    [$.array_type],
    [$.array_type, $.expression],
    [$.array_type, $.array_literal],
    [$.expression_statement, $.parenthesized_expression],
  ],

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

    _top_level_declaration: $ => choice(
      $.import_declaration,
      $.function_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),
    ),

    function_declaration: $ => seq(
      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(
      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(
      field('name', $.identifier),
      optional(field('type', $.type)),
      '::',
      repeat($._newline),
      field('value', $._value),
    ),

    global_variable_declaration: $ => seq(
      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: $ => seq(
      field('name', $.identifier),
      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), field('name', choice($.identifier, $.sink)))),
      field('type', $.type),
    ),

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

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

    named_type: $ => prec.right(seq(
      $.qualified_identifier,
      optional($.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', 'anyopaque', 'bool', 'int', '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('=', '+=', '-=', '*=', '/=')),
      repeat($._newline),
      field('right', $._value),
    ),

    expression_statement: $ => $.expression,

    return_statement: $ => seq(
      'return',
      repeat($._newline),
      field('value', $._value),
    ),

    yield_statement: $ => seq(
      'yield',
      repeat($._newline),
      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: $ => seq('defer', repeat($._newline), field('body', $.statement)),

    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(
      'for',
      repeat($._newline),
      field('iterable', $.expression),
      repeat($._newline),
      '|',
      optional('@'),
      field('item', $.identifier),
      optional(seq(',', field('index', $.identifier))),
      '|',
      repeat($._newline),
      optional(seq(field('label', $.identifier), ':', repeat($._newline))),
      field('body', $.block),
    ),

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

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

    match_capture: $ => seq('|', optional('@'), field('name', 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,
      $.call_expression,
      $.index_expression,
      $.slice_expression,
      $.postfix_expression,
      $.struct_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,
      $.none,
      $.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.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', $.identifier),
    )),

    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),
          $.field_initializer,
          repeat(seq(repeat($._newline), ',', repeat($._newline), $.field_initializer)),
          optional(seq(repeat($._newline), ',')),
          repeat($._newline),
        ),
      ),
      '}',
    ),

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

    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', $.identifier),
      '=',
      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'),
    none: _ => 'none',
    undefined: _ => 'undefined',
    sink: _ => '_',

    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)),
  );
}