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

  word: $ => $.identifier,

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

  rules: {
    root: $ =>
      prec(1, repeat1($.statement)),

    _comment: $ => /[#][^#\n]*[#|\n]/,

    statement: $ =>
      prec.left(
        seq(
          $.statement_kind,
          optional(';'),
        ),
      ),

    break: $ =>
      prec.left(
        seq(
          'break',
          optional($.statement),
        ),    
      ),

    return: $ =>
      prec.left(
        seq(
          'return',
          optional($.statement),
        ),      
      ),

    statement_kind: $ =>
      prec.left(
        choice(
          $.assignment,
          $.block,
          $.break,
          $.expression,
          $.for,
          $.if_else,
          $.index_assignment,
          $.loop_node,
          $.match,
          $.pipe,
          $.return,
          $.while,
          $.type_definition,
        ),
      ),

    expression: $ =>
      choice(
        $._expression_kind,
        seq(
          '(',
          $._expression_kind,
          ')',
        ),
      ),

    _expression_kind: $ =>
      prec.right(
        choice(
          $.as_node,
          $.function_call,
          $.identifier,
          $.index,
          $.logic,
          $.math,
          $.value,
          $.command,
        ),
      ),

    _expression_list: $ =>
      repeat1(
        prec.right(
          seq(
            $.expression,
            optional(','),
          ),
        ),
      ),

    as_node: $ =>
      seq($.expression, 'as', $.type),

    pipe: $ =>
      prec(
        1,
        seq(
          choice(
            $.command,
            $.function_call,
          ),
          '|',
          choice(
            $.command,
            $.pipe,
            $.function_call,
          ),
        ),
      ),

    command: $ =>
      prec.right(
        seq(
          '^',
          $.command_text,
          repeat($.command_argument),
        ),
      ),

    command_text: $ => /[^\s;]+/,

    command_argument: $ =>
      choice(
        /[^^|;\s]+/,
        /("[^"]*?")|('[^']*?')|(`[^`]*?`)/,
      ),

    block: $ =>
      seq(
        optional('async'),
        '{',
        repeat($.statement),
        '}',
      ),

    identifier: $ =>
      /[_a-zA-Z]+[_a-zA-Z0-9]*[_a-zA-Z]?/,

    value: $ =>
      choice(
        $.anonymous_function,
        $.integer,
        $.float,
        $.string,
        $.boolean,
        $.list,
        $.map,
        $.range,
        $.struct_instance,
        $.enum_instance,
      ),

    range: $ => /\d+[.]{2}\d+/,

    integer: $ => /[-]?\d+/,

    float: $ =>
      choice(
        /[-|+]?\d*[.][\d|e|-]*/,
        'Infinity',
        'infinity',
        'NaN',
        'nan',
      ),

    string: $ =>
      /("[^"]*?")|('[^']*?')|(`[^`]*?`)/,

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

    list: $ =>
      seq(
        '[',
        repeat(
          prec.left(
            seq(
              $.expression,
              optional(','),
            ),
          ),
        ),
        ']',
      ),

    map: $ =>
      prec(
        1,
        seq(
          '{',
          repeat(
            seq(
              $.identifier,
              optional(
                $.type_specification,
              ),
              '=',
              $.statement,
              optional(','),
            ),
          ),
          '}',
        ),
      ),

    index: $ =>
      prec.left(
        1,
        seq(
          $.index_expression,
          ':',
          $.index_expression,
        ),
      ),

    index_expression: $ =>
      prec(
        1,
        choice(
          seq(
            '(',
            $.function_call,
            ')',
          ),
          $.identifier,
          $.index,
          $.value,
          $.range,
        ),
      ),

    math: $ =>
      prec.left(
        seq(
          $.expression,
          $.math_operator,
          $.expression,
        ),
      ),

    math_operator: $ =>
      choice('+', '-', '*', '/', '%'),

    logic: $ =>
      prec.left(
        seq(
          $.expression,
          $.logic_operator,
          $.expression,
        ),
      ),

    logic_operator: $ =>
      prec.left(
        choice(
          '==',
          '!=',
          '&&',
          '||',
          '>',
          '<',
          '>=',
          '<=',
        ),
      ),

    assignment: $ =>
      seq(
        $.identifier,
        optional($.type_specification),
        $.assignment_operator,
        $.statement,
      ),

    index_assignment: $ =>
      seq(
        $.index,
        $.assignment_operator,
        $.statement,
      ),

    assignment_operator: $ =>
      prec.right(
        choice('=', '+=', '-='),
      ),

    if_else: $ =>
      prec.right(
        seq(
          $.if,
          repeat($.else_if),
          optional($.else),
        ),
      ),

    if: $ =>
      seq('if', $.expression, $.block),

    else_if: $ =>
      seq(
        'else if',
        $.expression,
        $.block,
      ),

    else: $ => seq('else', $.block),

    match: $ =>
      prec.right(
        seq(
          'match',
          $.expression,
          '{',
          repeat1(
            seq(
              $.match_pattern,
              '->',
              $.statement,
              optional(','),
            ),
          ),
          '}',
        ),
      ),

    match_pattern: $ =>
      choice(
        $.enum_pattern,
        $.value,
        '*',
      ),

    enum_pattern: $ =>
      prec(
        1,
        seq(
          $.identifier,
          '::',
          $.identifier,
          optional(
            seq('(', $.identifier, ')'),
          ),
        ),
      ),

    loop_node: $ =>
      seq('loop', $.block),

    while: $ =>
      seq(
        'while',
        $.expression,
        $.block,
      ),

    for: $ =>
      seq(
        choice('for', 'async for'),
        $.identifier,
        'in',
        $.expression,
        $.block,
      ),

    type_specification: $ =>
      seq('<', $.type, '>'),

    type: $ =>
      prec.right(
        choice(
          'any',
          'bool',
          'collection',
          'float',
          'int',
          'list',
          'map',
          'none',
          'num',
          'str',
          // Custom type
          $.identifier,

          // Custom type with arguments
          seq(
            $.identifier,
            $.type_arguments,
          ),

          // Map with exact fields
          seq(
            '{',
            repeat1(
              seq(
                $.identifier,
                $.type_specification,
              ),
            ),
            '}',
          ),

          // List of
          seq('list', '<', $.type, '>'),

          // Exact list
          seq(
            '[',
            repeat(
              seq(
                $.type,
                optional(','),
              ),
            ),
            ']',
          ),

          // Function
          seq(
            '(',
            repeat(
              seq(
                $.type,
                optional(','),
              ),
            ),
            ')',
            optional(seq('->', $.type)),
          ),
        ),
      ),

    anonymous_function: $ =>
      seq(
        '(',
        repeat(
          seq(
            $.identifier,
            $.type_specification,
            optional(','),
          ),
        ),
        ')',
        $.type_specification,
        $.block,
      ),

    function_expression: $ =>
      choice(
        $._function_expression_kind,
        seq(
          '(',
          $._function_expression_kind,
          ')',
        ),
      ),

    _function_expression_kind: $ =>
      prec(
        2,
        choice(
          $.function_call,
          $.identifier,
          $.index,
          $.value,
        ),
      ),

    function_call: $ =>
      prec.right(
        seq(
          $.function_expression,
          '(',
          optional($._expression_list),
          ')',
        ),
      ),

    type_definition: $ =>
      choice(
        $.enum_definition,
        $.struct_definition,
      ),

    type_arguments: $ =>
      seq(
        '<',
        repeat1(
          seq($.type, optional(',')),
        ),
        '>',
      ),

    enum_definition: $ =>
      prec.right(
        seq(
          'enum',
          $.identifier,
          optional($.type_arguments),
          repeat(
            seq(
              '{',
              repeat1(
                seq(
                  $.identifier,
                  optional(
                    $.type_arguments,
                  ),
                  optional(','),
                ),
              ),
              '}',
            ),
          ),
        ),
      ),

    enum_instance: $ =>
      prec.right(
        seq(
          $.identifier,
          '::',
          $.identifier,
          optional(
            seq('(', $.expression, ')'),
          ),
        ),
      ),

    struct_definition: $ =>
      seq(
        'struct',
        $.identifier,
        '{',
        repeat(
          choice(
            seq(
              $.identifier,
              $.type_specification,
            ),
            seq(
              $.identifier,
              '=',
              $.statement,
            ),
            seq(
              $.identifier,
              $.type_specification,
              '=',
              $.statement,
            ),
          ),
        ),
        '}',
      ),

    struct_instance: $ =>
      seq($.identifier, '::', $.map),
  },
});