// Tree-sitter grammar for the Lux programming language
module.exports = grammar({
  name: 'lux',

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

  conflicts: $ => [
    [$.primary_expression, $.pattern],
    [$.type_expression, $.primary_expression],
  ],

  word: $ => $.identifier,

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

    _declaration: $ => choice(
      $.function_declaration,
      $.let_declaration,
      $.type_declaration,
      $.effect_declaration,
      $.handler_declaration,
      $.trait_declaration,
      $.impl_declaration,
      $.import_declaration,
      $.export_declaration,
    ),

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

    // Function declaration
    function_declaration: $ => seq(
      'fn',
      field('name', $.identifier),
      optional($.type_parameters),
      $.parameter_list,
      ':',
      field('return_type', $.type_expression),
      optional($.effect_clause),
      optional($.property_clause),
      '=',
      field('body', $._expression),
    ),

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

    parameter: $ => seq(
      field('name', $.identifier),
      ':',
      field('type', $.type_expression),
    ),

    type_parameters: $ => seq(
      '<',
      $.identifier,
      repeat(seq(',', $.identifier)),
      optional(','),
      '>',
    ),

    effect_clause: $ => seq(
      'with',
      '{',
      optional(seq(
        $.identifier,
        repeat(seq(',', $.identifier)),
        optional(','),
      )),
      '}',
    ),

    property_clause: $ => seq(
      'is',
      $.identifier,
      repeat(seq(',', $.identifier)),
    ),

    // Let declaration
    let_declaration: $ => seq(
      'let',
      field('name', $.identifier),
      optional(seq(':', field('type', $.type_expression))),
      '=',
      field('value', $._expression),
    ),

    // Type declaration
    type_declaration: $ => seq(
      'type',
      field('name', $.identifier),
      optional($.type_parameters),
      '=',
      field('definition', $.type_definition),
    ),

    type_definition: $ => choice(
      $.enum_definition,
      $.record_type,
      $.type_expression,
    ),

    enum_definition: $ => seq(
      repeat1(seq('|', $.variant)),
    ),

    variant: $ => seq(
      field('name', $.identifier),
      optional(choice(
        $.tuple_fields,
        $.record_fields,
      )),
    ),

    tuple_fields: $ => seq(
      '(',
      $.type_expression,
      repeat(seq(',', $.type_expression)),
      optional(','),
      ')',
    ),

    record_fields: $ => seq(
      '{',
      $.record_field,
      repeat(seq(',', $.record_field)),
      optional(','),
      '}',
    ),

    record_field: $ => seq(
      field('name', $.identifier),
      ':',
      field('type', $.type_expression),
    ),

    // Effect declaration
    effect_declaration: $ => seq(
      'effect',
      field('name', $.identifier),
      optional($.type_parameters),
      '{',
      repeat($.effect_operation),
      '}',
    ),

    effect_operation: $ => seq(
      'fn',
      field('name', $.identifier),
      $.parameter_list,
      ':',
      field('return_type', $.type_expression),
    ),

    // Handler declaration
    handler_declaration: $ => seq(
      'handler',
      field('name', $.identifier),
      ':',
      field('effect', $.identifier),
      '{',
      repeat($.handler_operation),
      '}',
    ),

    handler_operation: $ => seq(
      'fn',
      field('name', $.identifier),
      '(',
      optional(seq(
        $.identifier,
        repeat(seq(',', $.identifier)),
        optional(','),
      )),
      ')',
      '=',
      field('body', $._expression),
    ),

    // Trait declaration
    trait_declaration: $ => seq(
      'trait',
      field('name', $.identifier),
      optional($.type_parameters),
      '{',
      repeat($.trait_method),
      '}',
    ),

    trait_method: $ => seq(
      'fn',
      field('name', $.identifier),
      $.parameter_list,
      ':',
      field('return_type', $.type_expression),
    ),

    // Impl declaration
    impl_declaration: $ => seq(
      'impl',
      field('trait', $.identifier),
      'for',
      field('type', $.type_expression),
      '{',
      repeat($.function_declaration),
      '}',
    ),

    // Import/Export
    import_declaration: $ => seq(
      'import',
      $.import_path,
      optional($.import_clause),
    ),

    import_path: $ => seq(
      $.identifier,
      repeat(seq('.', $.identifier)),
    ),

    import_clause: $ => seq(
      '{',
      $.identifier,
      repeat(seq(',', $.identifier)),
      optional(','),
      '}',
    ),

    export_declaration: $ => seq(
      'export',
      '{',
      $.identifier,
      repeat(seq(',', $.identifier)),
      optional(','),
      '}',
    ),

    // Type expressions
    type_expression: $ => choice(
      $.identifier,
      $.generic_type,
      $.function_type,
      $.tuple_type,
      $.record_type,
      seq('(', $.type_expression, ')'),
    ),

    generic_type: $ => seq(
      $.identifier,
      '<',
      $.type_expression,
      repeat(seq(',', $.type_expression)),
      optional(','),
      '>',
    ),

    function_type: $ => seq(
      'fn',
      '(',
      optional(seq(
        $.type_expression,
        repeat(seq(',', $.type_expression)),
        optional(','),
      )),
      ')',
      ':',
      $.type_expression,
      optional($.effect_clause),
    ),

    tuple_type: $ => seq(
      '(',
      $.type_expression,
      ',',
      $.type_expression,
      repeat(seq(',', $.type_expression)),
      optional(','),
      ')',
    ),

    record_type: $ => seq(
      '{',
      optional(seq(
        $.record_field,
        repeat(seq(',', $.record_field)),
        optional(','),
      )),
      '}',
    ),

    // Expressions
    _expression: $ => choice(
      $.primary_expression,
      $.binary_expression,
      $.unary_expression,
      $.call_expression,
      $.member_expression,
      $.index_expression,
      $.if_expression,
      $.match_expression,
      $.block_expression,
      $.lambda_expression,
      $.run_expression,
      $.resume_expression,
      $.pipe_expression,
    ),

    primary_expression: $ => choice(
      $.identifier,
      $.number,
      $.string,
      $.char,
      $.boolean,
      $.unit,
      $.list_expression,
      $.tuple_expression,
      $.record_expression,
      seq('(', $._expression, ')'),
    ),

    // Literals
    number: $ => choice(
      $.integer,
      $.float,
    ),

    integer: $ => /\d+/,

    float: $ => /\d+\.\d+/,

    string: $ => choice(
      $.simple_string,
      $.interpolated_string,
    ),

    simple_string: $ => seq(
      '"',
      repeat(choice(
        $.string_content,
        $.escape_sequence,
      )),
      '"',
    ),

    interpolated_string: $ => seq(
      '"',
      repeat(choice(
        $.string_content,
        $.escape_sequence,
        $.interpolation,
      )),
      '"',
    ),

    string_content: $ => /[^"\\{]+/,

    escape_sequence: $ => /\\[nrt\\'"0{}\}]/,

    interpolation: $ => seq(
      '{',
      $._expression,
      '}',
    ),

    char: $ => seq(
      "'",
      choice(
        /[^'\\]/,
        $.escape_sequence,
      ),
      "'",
    ),

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

    unit: $ => seq('(', ')'),

    list_expression: $ => seq(
      '[',
      optional(seq(
        $._expression,
        repeat(seq(',', $._expression)),
        optional(','),
      )),
      ']',
    ),

    tuple_expression: $ => seq(
      '(',
      $._expression,
      ',',
      $._expression,
      repeat(seq(',', $._expression)),
      optional(','),
      ')',
    ),

    record_expression: $ => seq(
      '{',
      optional(seq(
        $.field_assignment,
        repeat(seq(',', $.field_assignment)),
        optional(','),
      )),
      '}',
    ),

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

    // Binary expressions
    binary_expression: $ => choice(
      prec.left(1, seq($._expression, '||', $._expression)),
      prec.left(2, seq($._expression, '&&', $._expression)),
      prec.left(3, seq($._expression, choice('==', '!='), $._expression)),
      prec.left(4, seq($._expression, choice('<', '>', '<=', '>='), $._expression)),
      prec.left(5, seq($._expression, choice('+', '-'), $._expression)),
      prec.left(6, seq($._expression, choice('*', '/', '%'), $._expression)),
    ),

    unary_expression: $ => prec(7, choice(
      seq('-', $._expression),
      seq('!', $._expression),
    )),

    // Pipe expression
    pipe_expression: $ => prec.left(0, seq(
      $._expression,
      '|>',
      $._expression,
    )),

    // Call expression
    call_expression: $ => prec(8, seq(
      field('function', $._expression),
      '(',
      optional(seq(
        $._expression,
        repeat(seq(',', $._expression)),
        optional(','),
      )),
      ')',
    )),

    // Member expression
    member_expression: $ => prec(9, seq(
      field('object', $._expression),
      '.',
      field('member', $.identifier),
    )),

    // Index expression
    index_expression: $ => prec(8, seq(
      field('object', $._expression),
      '[',
      field('index', $._expression),
      ']',
    )),

    // If expression
    if_expression: $ => prec.right(seq(
      'if',
      field('condition', $._expression),
      'then',
      field('then', $._expression),
      'else',
      field('else', $._expression),
    )),

    // Match expression
    match_expression: $ => seq(
      'match',
      field('value', $._expression),
      '{',
      repeat1($.match_arm),
      '}',
    ),

    match_arm: $ => seq(
      field('pattern', $.pattern),
      '=>',
      field('body', $._expression),
      optional(','),
    ),

    pattern: $ => choice(
      $.identifier,
      $.number,
      $.string,
      $.boolean,
      '_',
      $.constructor_pattern,
      $.tuple_pattern,
      $.record_pattern,
    ),

    constructor_pattern: $ => seq(
      field('name', $.identifier),
      optional(seq(
        '(',
        $.pattern,
        repeat(seq(',', $.pattern)),
        optional(','),
        ')',
      )),
    ),

    tuple_pattern: $ => seq(
      '(',
      $.pattern,
      ',',
      $.pattern,
      repeat(seq(',', $.pattern)),
      optional(','),
      ')',
    ),

    record_pattern: $ => seq(
      '{',
      $.field_pattern,
      repeat(seq(',', $.field_pattern)),
      optional(','),
      '}',
    ),

    field_pattern: $ => seq(
      field('name', $.identifier),
      optional(seq(':', field('pattern', $.pattern))),
    ),

    // Block expression
    block_expression: $ => seq(
      '{',
      repeat($.statement),
      optional($._expression),
      '}',
    ),

    statement: $ => seq(
      $.let_declaration,
    ),

    // Lambda expression
    lambda_expression: $ => seq(
      'fn',
      $.parameter_list,
      optional(seq(':', $.type_expression)),
      '=>',
      $._expression,
    ),

    // Run expression
    run_expression: $ => seq(
      'run',
      field('expression', $._expression),
      'with',
      '{',
      optional(seq(
        $.handler_binding,
        repeat(seq(',', $.handler_binding)),
        optional(','),
      )),
      '}',
    ),

    handler_binding: $ => seq(
      field('effect', $.identifier),
      '=',
      field('handler', $.identifier),
    ),

    // Resume expression
    resume_expression: $ => seq(
      'resume',
      '(',
      $._expression,
      ')',
    ),

    // Identifier
    identifier: $ => /[a-zA-Z_][a-zA-Z0-9_]*/,
  }
});