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

  rules: {
    source_file: $ => repeat($._item),

    _item: $ => choice(
      $.command,
      $.comment,
      $._whitespace
    ),

    // VM commands
    command: $ => choice(
      $.arithmetic_command,
      $.memory_command,
      $.program_flow_command,
      $.function_command
    ),

    // Arithmetic and logical commands
    arithmetic_command: $ => choice(
      'add', 'sub', 'neg',
      'eq', 'gt', 'lt',
      'and', 'or', 'not'
    ),

    // Memory access commands
    memory_command: $ => choice(
      $.push_command,
      $.pop_command
    ),

    push_command: $ => seq(
      'push',
      $.memory_segment,
      $.index
    ),

    pop_command: $ => seq(
      'pop',
      $.memory_segment,
      $.index
    ),

    // Memory segments
    memory_segment: $ => choice(
      'argument', 'local', 'static', 'constant',
      'this', 'that', 'pointer', 'temp'
    ),

    // Program flow commands
    program_flow_command: $ => choice(
      $.label_command,
      $.goto_command,
      $.if_goto_command
    ),

    label_command: $ => seq(
      'label',
      $.label_name
    ),

    goto_command: $ => seq(
      'goto',
      $.label_name
    ),

    if_goto_command: $ => seq(
      'if-goto',
      $.label_name
    ),

    // Function calling commands
    function_command: $ => choice(
      $.function_declaration,
      $.call_command,
      $.return_command
    ),

    function_declaration: $ => seq(
      'function',
      $.function_name,
      $.local_vars_count
    ),

    call_command: $ => seq(
      'call',
      $.function_name,
      $.args_count
    ),

    return_command: $ => 'return',

    // Identifiers and numbers
    index: $ => /\d+/,
    local_vars_count: $ => /\d+/,
    args_count: $ => /\d+/,

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

    // Comments
    comment: $ => token(seq('//', /.*/)),

    // Whitespace
    _whitespace: $ => /\s+/
  },

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