const PREC = {
  ASSIGNMENT: 1,
  LOGICAL: 2,    // or, &, xor
  RELATIONAL: 3, // =, <>, <, >, <=, >=
  SHIFT: 4,      // <<, >>
  ADD: 5,        // +, -
  MULTIPLY: 6,   // *, /, %
  UNARY: 7,      // @, ~, -, +
  CALL: 8,       // function calls
  POSTFIX: 9,    // ., [], [to], ^
};

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

function sep1(rule, separator) {
  return seq(rule, repeat(seq(separator, rule)));
}

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

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

  word: $ => $.identifier,

  conflicts: $ => [
    [$._statement, $._expression],
  ],

  rules: {
    source_file: $ => seq(
      choice(
        $.module_declaration,
        $.program_declaration,
      ),
      optional($.trailing_content),
    ),

    trailing_content: $ => token(seq(/[^\s]/, repeat(/[\s\S]/))),

    comment: $ => token(choice(
      seq('(*', /[^*]*\*+([^)*][^*]*\*+)*/, ')'),
      seq('--', /[^\n]*/),
    )),

    identifier: $ => /[A-Za-z_][A-Za-z0-9_]*/,

    trait: $ => token(seq('#', /[A-Za-z_][A-Za-z0-9_]*/)),

    module_declaration: $ => seq(
      'module',
      repeat($._declaration),
      'end',
      '.',
    ),

    program_declaration: $ => seq(
      'program',
      '(',
      optional(field('parameters', commaSep1($.identifier))),
      ')',
      choice(
        seq(
          repeat1($._declaration),
          optional(seq('begin', optional($._statements))),
        ),
        seq('begin', optional($._statements)),
        optional($._statements),
      ),
      'end',
      '.',
    ),

    _declaration: $ => choice(
      $.import_declaration,
      $.type_section,
      $.const_section,
      $.var_section,
      $.threadvar_section,
      $.procedure_declaration,
    ),

    import_declaration: $ => seq(
      'import',
      commaSep1($.import_path),
    ),

    import_path: $ => seq(
      $.identifier,
      repeat(seq('.', $.identifier)),
      optional(field('exported', '*')),
    ),

    type_section: $ => seq(
      'type',
      repeat(seq($.type_declaration, optional(';'))),
    ),

    type_declaration: $ => seq(
      field('name', $.identifier_definition),
      '=',
      field('type', choice($._type_expression, 'extern')),
    ),

    const_section: $ => seq(
      'const',
      repeat(seq($.const_declaration, optional(';'))),
    ),

    const_declaration: $ => seq(
      field('name', $.identifier_definition),
      '=',
      field('value', $._expression),
    ),

    identifier_definition: $ => seq(
      field('name', $.identifier),
      optional(field('type_parameters', $.type_parameters)),
      optional(field('exported', '*')),
      repeat($.attribute),
    ),

    attribute: $ => seq(
      '#{',
      field('name', $.identifier),
      optional(field('arguments', commaSep1($._expression))),
      '}',
    ),

    type_parameters: $ => seq(
      '#[',
      commaSep1($.identifier),
      ']',
    ),

    var_section: $ => seq(
      'var',
      repeat(seq($.variable_declaration, optional(';'))),
    ),

    threadvar_section: $ => seq(
      'threadvar',
      repeat(seq($.variable_declaration, optional(';'))),
    ),

    variable_declaration: $ => seq(
      field('names', commaSep1($.identifier_definition)),
      ':',
      field('type', $._type_expression),
      optional(seq(
        ':=',
        field('value', choice($._expression, 'extern')),
      )),
    ),

    procedure_declaration: $ => seq(
      'proc',
      field('name', $.identifier_definition),
      field('parameters', $.parameter_list),
      repeat($.attribute),
      optional(seq('->', field('return_type', $._type_expression))),
      field('body', choice(
        $.procedure_body,
        'extern',
      )),
    ),

    parameter_list: $ => seq(
      '(',
      optional(seq(
        $.field_declaration,
        repeat(seq(';', $.field_declaration)),
        optional(';'),
      )),
      ')',
    ),

    field_identifier: $ => seq(
      field('name', $.identifier),
      repeat($.attribute),
    ),

    field_declaration: $ => seq(
      field('names', commaSep1($.field_identifier)),
      ':',
      field('type', $._type_expression),
    ),

    procedure_body: $ => seq(
      optional($.block),
      $.return_clause,
    ),

    return_clause: $ => seq(
      'return',
      optional(field('value', $._expression)),
    ),

    _block_declaration: $ => choice(
      $.const_section,
      $.var_section,
    ),

    block: $ => choice(
      seq(
        repeat1($._block_declaration),
        optional(seq('begin', optional($._statements))),
      ),
      seq(
        'begin',
        optional($._statements),
      ),
      $._statements,
    ),

    _statements: $ => repeat1(choice($._statement, ';')),

    _statement: $ => choice(
      $.assignment_statement,
      $.call_expression,
      $.if_statement,
      $.while_statement,
      $.repeat_statement,
      $.for_statement,
      $.defer_statement,
      $.block_statement,
      $.break_statement,
      $.case_statement,
    ),

    assignment_statement: $ => prec.right(PREC.ASSIGNMENT, seq(
      field('left', $._expression),
      ':=',
      field('right', $._expression),
    )),

    if_statement: $ => seq(
      'if',
      field('condition', $._expression),
      'then',
      field('consequence', optional($.block)),
      repeat($.elsif_clause),
      optional($.else_clause),
      'end',
    ),

    elsif_clause: $ => seq(
      'elsif',
      field('condition', $._expression),
      'then',
      field('consequence', optional($.block)),
    ),

    else_clause: $ => seq(
      'else',
      field('body', optional($.block)),
    ),

    while_statement: $ => seq(
      'while',
      field('condition', $._expression),
      'do',
      field('body', optional($.block)),
      repeat($.while_elsif_clause),
      'end',
    ),

    while_elsif_clause: $ => seq(
      'elsif',
      field('condition', $._expression),
      'do',
      field('body', optional($.block)),
    ),

    repeat_statement: $ => seq(
      'repeat',
      field('body', optional($.block)),
      'until',
      field('condition', $._expression),
    ),

    for_statement: $ => seq(
      'for',
      field('variable', $.identifier),
      'of',
      field('iterable', $._expression),
      optional(seq('with', field('counter', $.identifier))),
      'do',
      field('body', optional($.block)),
      'end',
    ),

    defer_statement: $ => seq(
      'defer',
      field('body', optional($.block)),
      'end',
    ),

    block_statement: $ => seq(
      'block',
      field('name', $.identifier),
      field('body', optional($.block)),
      'end',
    ),

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

    case_statement: $ => seq(
      'case',
      field('value', $._expression),
      'of',
      optional('|'),
      sep1($.case_arm, '|'),
      optional($.else_clause),
      'end',
    ),

    case_arm: $ => seq(
      field('patterns', commaSep1($._expression)),
      ':',
      field('body', optional($.block)),
    ),

    // Type expressions
    _type_expression: $ => choice(
      $.array_type,
      $.slice_type,
      $.const_type,
      $.pointer_type,
      $.record_type,
      $.procedure_type,
      $.enum_type,
      $.generic_type,
      $.identifier,
    ),

    array_type: $ => prec(1, seq(
      '[',
      field('length', $._expression),
      ']',
      field('element_type', $._type_expression),
    )),

    slice_type: $ => prec(1, seq(
      '[',
      ']',
      field('element_type', $._type_expression),
    )),

    const_type: $ => seq(
      'const',
      field('type', $._type_expression),
    ),

    pointer_type: $ => prec.right(seq(
      '^',
      field('type', $._type_expression),
    )),

    record_type: $ => seq(
      'record',
      optional(seq('(', field('base', choice($.identifier, $.generic_type)), ')')),
      repeat($.attribute),
      repeat(seq($.field_declaration, optional(';'))),
      'end',
    ),

    procedure_type: $ => seq(
      'proc',
      field('parameters', $.parameter_list),
      repeat($.attribute),
      optional(seq('->', field('return_type', $._type_expression))),
    ),

    enum_type: $ => seq(
      '(',
      commaSep1(field('member', $.enum_member)),
      ')',
    ),

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

    generic_type: $ => seq(
      field('name', $.identifier),
      field('type_arguments', $.type_arguments),
    ),

    type_arguments: $ => seq(
      '#[',
      commaSep1($._type_expression),
      ']',
    ),

    // Expressions
    _expression: $ => choice(
      $.literal,
      $.identifier,
      $.generic_type,
      $.unary_expression,
      $.binary_expression,
      $.call_expression,
      $.index_expression,
      $.slice_expression,
      $.field_expression,
      $.dereference_expression,
      $.cast_expression,
      $.traits_expression,
      $.record_constructor,
      $.array_constructor,
      $.parenthesized_expression,
    ),

    unary_expression: $ => prec.right(PREC.UNARY, seq(
      field('operator', choice('@', '~', '-', '+')),
      field('argument', $._expression),
    )),

    binary_expression: $ => choice(
      ...[
        ['or', PREC.LOGICAL],
        ['&', PREC.LOGICAL],
        ['xor', PREC.LOGICAL],
        ['=', PREC.RELATIONAL],
        ['<>', PREC.RELATIONAL],
        ['<', PREC.RELATIONAL],
        ['>', PREC.RELATIONAL],
        ['<=', PREC.RELATIONAL],
        ['>=', PREC.RELATIONAL],
        ['<<', PREC.SHIFT],
        ['>>', PREC.SHIFT],
        ['+', PREC.ADD],
        ['-', PREC.ADD],
        ['*', PREC.MULTIPLY],
        ['/', PREC.MULTIPLY],
        ['%', PREC.MULTIPLY],
      ].map(([operator, precedence]) =>
        prec.left(precedence, seq(
          field('left', $._expression),
          field('operator', operator),
          field('right', $._expression),
        ))
      )
    ),

    call_expression: $ => prec(PREC.CALL, seq(
      field('function', $._expression),
      field('arguments', $.argument_list),
    )),

    argument_list: $ => seq(
      '(',
      optional(commaSep1($._expression)),
      ')',
    ),

    index_expression: $ => prec(PREC.POSTFIX, seq(
      field('array', $._expression),
      '[',
      field('index', $._expression),
      ']',
    )),

    slice_expression: $ => prec(PREC.POSTFIX, seq(
      field('array', $._expression),
      '[',
      field('start', $._expression),
      'to',
      field('end', $._expression),
      ']',
    )),

    field_expression: $ => prec(PREC.POSTFIX, seq(
      field('operand', $._expression),
      '.',
      field('field', $.identifier),
    )),

    dereference_expression: $ => prec(PREC.POSTFIX, seq(
      field('operand', $._expression),
      '^',
    )),

    cast_expression: $ => seq(
      'cast',
      '(',
      field('value', $._expression),
      ':',
      field('type', $._type_expression),
      ')',
    ),

    traits_expression: $ => seq(
      field('trait', $.trait),
      '(',
      field('arguments', commaSep1($._type_expression)),
      ')',
    ),

    record_constructor: $ => seq(
      '{',
      field('type', $._type_expression),
      optional(seq(';', commaSep1($.field_initializer))),
      '}',
    ),

    field_initializer: $ => seq(
      field('name', $.identifier),
      ':',
      field('value', $._expression),
    ),

    array_constructor: $ => seq(
      '[',
      commaSep1($._expression),
      ']',
    ),

    parenthesized_expression: $ => seq(
      '(',
      $._expression,
      ')',
    ),

    // Literals
    literal: $ => choice(
      $.integer,
      $.float,
      $.string,
      $.byte,
      $.character,
      $.boolean,
    ),

    boolean: $ => choice('true', 'false'),

    integer: $ => token(choice(
      seq(/0[xX][0-9a-fA-F]+/, optional(choice('u8', 'u16', 'u32', 'u64'))),
      seq(/0[bB][01]+/, optional(choice('u8', 'u16', 'u32', 'u64'))),
      seq(/[0-9]+/, optional(choice('u', 'u8', 'u16', 'u32', 'u64', 'i8', 'i16', 'i32', 'i64'))),
    )),

    float: $ => token(seq(
      /[0-9]+/,
      choice(
        seq('.', /[0-9]+/, optional(/[eE][+-]?[0-9]+/)),
        /[eE][+-]?[0-9]+/,
      ),
      optional('f'),
    )),

    string: $ => seq(
      '"',
      repeat(choice(
        token.immediate(prec(1, /[^"\\\n]+/)),
        $.escape_sequence,
        '\\\n',
      )),
      '"',
    ),

    byte: $ => seq(
      "'",
      repeat1(choice(
        token.immediate(prec(1, /[^'\\\n]+/)),
        $.escape_sequence,
        '\\\n',
      )),
      "'",
    ),

    character: $ => seq(
      '`',
      repeat1(choice(
        token.immediate(prec(1, /[^`\\\n]+/)),
        $.escape_sequence,
        '\\\n',
      )),
      '`',
    ),

    escape_sequence: $ => token.immediate(seq(
      '\\',
      choice(
        /[^0-9]/,
        /[0-9]{1,3}/,
      ),
    )),
  },
});