from nodes import *
import tokens

@with_lexer(tokens.lkt_lexer)
@with_unparsers
grammar lkt_grammar {
    @main_rule
    main_rule <- LangkitRoot(/ module_doc imports decls @Termination)
    id <- Id(@Identifier)
    ref_id <- RefId(@Identifier)
    type_ref_id <- ref_id |> when(Id.is_type_name)
    module_id <- ModuleId(@Identifier) |> when(Id.is_not_type_name)
    def_id <- DefId(@Identifier)
    doc <- ?string_lit
    module_doc <- ?ModuleDocStringLit(
        list+(ModuleDocStringLine(@ModuleDocStringLine))
    )
    imported_names <- list+(
        ImportedName(
            ImportedId(@Identifier)
            /
            ?pick(@Identifier("as") / def_id)
        ),
        ","
    )
    import_clause <- or(
        | Import("import" / imported_names)
        | ImportFrom("from" module_id "import" imported_names)
        | ImportAllFrom("from" / module_id "import" "*")
    )
    imports <- list*(import_clause)
    lexer_decl <- LexerDecl(
        "lexer"
        /
        def_id
        "{"
        list*(or(lexer_case_rule | decl)).dont_skip("}")
        "}"
    )
    grammar_decl <- GrammarDecl("grammar" / def_id "{" decls "}")
    grammar_rule <- GrammarRuleDecl(def_id "<-" grammar_expr)
    lexer_case_rule <- LexerCaseRule(
        "match" / grammar_primary "{" list*(lexer_case_alt).dont_skip("}") "}"
    )
    lexer_case_alt <- or(
        | LexerCaseRuleCondAlt(
            "if"
            /
            @Identifier("previous_token")
            "is"
            list+(ref_id, "|")
            "then"
            lexer_case_send
        )
        | LexerCaseRuleDefaultAlt("else" / lexer_case_send)
        | skip(ErrorLexerCaseRuleAlt)
    )
    lexer_case_send <- LexerCaseRuleSend(
        @Identifier("send") "(" ref_id "," num_lit ")"
    )
    grammar_primary <- or(
        | grammar_pick
        | grammar_list_expr
        | token_literal
        | token_no_case_literal
        | token_pattern
        | grammar_cut
        | grammar_skip
        | grammar_null
        | grammar_token
        | grammar_stopcut
        | parse_node_expr
        | grammar_opt
        | grammar_opt_error
        | grammar_or_expr
        | grammar_rule_ref
        | grammar_discard_expr
        | skip(ErrorGrammarExpr)
    )
    grammar_expr <- or(
        | GrammarDontSkip(
            grammar_expr "." @Identifier("dont_skip") "(" grammar_expr ")"
        )
        | GrammarPredicate(
            grammar_expr "|>" "when" "(" basic_name ")"
        )
        | grammar_primary
    )
    grammar_pick <- GrammarPick(
        @Identifier("pick") "(" list+(grammar_expr).dont_skip(")") ")"
    )
    grammar_implicit_pick <- GrammarImplicitPick(list+(grammar_expr))
    grammar_opt <- or(
        | GrammarOptGroup("?" "(" / list*(grammar_expr).dont_skip(")") ")")
        | GrammarOpt("?" / grammar_expr)
    )
    grammar_opt_error <- or(
        | GrammarOptErrorGroup(
            "!" "(" / list*(grammar_expr).dont_skip(")") ")"
        )
        | GrammarOptError("!" / grammar_expr)
    )
    grammar_cut <- GrammarCut("/")
    grammar_stopcut <- GrammarStopCut(
        @Identifier("stop_cut") "(" grammar_expr ")"
    )
    grammar_or_expr <- GrammarOrExpr(
        "or"
        "("
        list+(
            list+(grammar_expr).dont_skip("|"), "|", allow_leading
        ).dont_skip(")")
        ")"
    )
    grammar_discard_expr <- GrammarDiscard("discard" "(" grammar_expr ")")
    token_literal <- TokenLit(@String)
    token_no_case_literal <- TokenNoCaseLit(
        @Identifier("no_case") "(" token_literal ")"
    )
    token_pattern <- or(
        | TokenPatternConcat(token_pattern "&" token_pattern_literal)
        | token_pattern_literal
    )
    token_pattern_literal <- TokenPatternLit(@PString)
    parse_node_expr <- ParseNodeExpr(
        type_ref
        "("
        /
        list*(grammar_expr).dont_skip(")")
        ")"
    )
    grammar_rule_ref <- GrammarRuleRef(ref_id)
    grammar_list_expr <- GrammarList(
        # Match either "list" (type inference will determine the list type) or
        # a specific list type.
        or(DefaultListTypeRef(@Identifier("list")) | type_ref)

        or(ListKind.One("+") | ListKind.Zero("*"))
        "("

        # Main list expr
        or(grammar_implicit_pick | grammar_expr).dont_skip(")").dont_skip(",")

        # Separator
        ?pick("," grammar_list_sep)
        ")"
    )
    grammar_list_sep <- GrammarListSep(grammar_expr ?pick("," id))
    grammar_skip <- GrammarSkip(@Identifier("skip") "(" type_ref ")")
    grammar_null <- GrammarNull("null" ?pick("(" type_ref ")"))
    grammar_token <- TokenRef("@" ref_id ?pick("(" token_literal ")"))
    type_decl <- or(
        | StructDecl(
            "struct"
            /
            def_id
            ?pick("implements" / type_list)
            "{"
            decl_block.dont_skip("}")
            "}"
        )
        | EnumClassDecl(
            "enum"
            "class"
            /
            def_id
            ?pick(":" type_ref)
            ?pick("implements" type_list)
            "{"
            list*(
                EnumClassCase(
                    "case"
                    list+(EnumClassAltDecl(def_id null(ASTList[TypeRef])), ",")
                )
            )
            decl_block.dont_skip("}")
            "}"
        )
        | ClassDecl(
            "class"
            /
            def_id
            ?pick(":" / type_ref)
            ?pick("implements" / type_list)
            "{"
            decl_block.dont_skip("}")
            "}"
        )
        | EnumTypeDecl(
            "enum"
            /
            def_id
            ?pick("implements" type_list)
            "{"
            "case"
            list+(enum_lit_decl, ",")
            decl_block.dont_skip("}")
            "}"
        )
        | TraitDecl(
            "trait"
            /
            def_id
            null(ASTList[TypeRef])
            "{"
            decl_block.dont_skip("}")
            "}"
        )
    )
    generic_decl <- GenericDecl(
        "generic"
        /
        "["
        GenericParamDeclList+(generic_param_type, ",")
        "]"
        bare_decl
    )
    generic_param_type <- FullDecl(
        doc
        list*(decl_annotation)
        GenericParamTypeDecl(
            ClassQualifier("class") def_id null(ASTList[TypeRef])
        )
    )
    enum_lit_decl <- EnumLitDecl(def_id)
    fun_decl <- FunDecl(
        "fun"
        /
        def_id
        "("
        fun_param_list
        ")"
        ":"
        type_ref
        ?pick("implements" type_member_ref)
        ?pick("=" expr)
    )
    lambda_param_decl <- LambdaParamDecl(
        def_id ?pick(":" type_ref) ?pick("=" expr)
    )
    fun_param_decl <- FunParamDecl(
        list*(decl_annotation) def_id ":" type_ref ?pick("=" expr)
    )
    fun_param_list <- list*(fun_param_decl, ",")
    lambda_param_list <- list*(lambda_param_decl, ",")
    field_decl <- FieldDecl(
        def_id
        ":"
        type_ref
        ?pick("implements" type_member_ref)
        ?pick("=" expr)
    )
    lexer_family_decl <- LexerFamilyDecl(
        @Identifier("family") / def_id "{" list*(decl).dont_skip("}") "}"
    )
    bare_decl <- or(
        | generic_decl
        | type_decl
        | fun_decl
        | lexer_decl
        | grammar_decl
        | field_decl
        | val_decl
        | env_spec_decl
        | lexer_family_decl
        | grammar_rule
        | dynvar_decl
        | skip(ErrorDecl)
    )
    decl <- FullDecl(doc list*(decl_annotation) bare_decl)
    type_member_ref <- DotExpr(type_ref_id null(NullCondQualifier)"." / ref_id)
    type_expr <- or(
        | DotExpr(module_id null(NullCondQualifier)"." / type_ref_id)
        | DotExpr(type_expr null(NullCondQualifier)"." / type_ref_id)
        | type_ref_id
    )
    type_ref <- or(
        | GenericTypeRef(type_expr "[" type_list "]")
        | SimpleTypeRef(type_expr)
        | FunctionTypeRef("(" list*(type_ref, ",") ")" "->" type_ref)
    )
    type_list <- list+(type_ref, ",")
    decls <- list*(decl).dont_skip("}")
    decl_block <- DeclBlock*(decl)
    val_decl <- ValDecl("val" / def_id ?pick(":" type_ref) "=" expr)
    dynvar_decl <- DynVarDecl("dynvar" def_id ":" type_ref)
    var_bind <- VarBind("bind" / ref_id "=" expr)
    env_spec_action <- CallExpr(RefId(@Identifier) "(" args ")")
    env_spec_decl <- EnvSpecDecl(
        DefId(@Identifier("env_spec")) "{" list*(env_spec_action) "}"
    )
    block <- BlockExpr(
        "{"
        /
        list+(
            or(
                | BlockExprClause(or(var_bind | decl) ";")
                | expr
                | skip(ErrorDecl)
            )
        ).dont_skip("}")
        "}"
    )
    pattern <- or(OrPattern(neg_pattern "|" pattern) | neg_pattern)
    neg_pattern <- or(
        | NotPattern("not" neg_pattern)
        | complex_pattern
    )
    pattern_binding <- BindingValDecl(def_id)
    complex_pattern <- or(
        | ComplexPattern(
                null(BindingValDecl)
                value_pattern
                ?pick("(" / list+(pattern_arg, ",") ")")
                ?pick("when" expr)
        )
        | ComplexPattern(
                ?pick(pattern_binding "@")
                value_pattern
                ?pick("(" / list+(pattern_arg, ",") ")")
                ?pick("when" expr)
        )
        | ComplexPattern(
                ?pick(pattern_binding "@")
                AnyTypePattern("*")
                "(" / list+(pattern_arg, ",") ")"
                ?pick("when" expr)
        )
        | RenamingComplexPattern(
                pattern_binding
                null(Pattern)
                null(ASTList[PatternDetail])
                ?pick("when" expr)
        )
    )
    value_pattern <- or(
        | TypePattern(type_ref)
        | NullPattern("null")
        | regex_pattern
        | bool_pattern
        | integer_pattern
        | list_pattern
        | ParenPattern("(" pattern ")")
    )
    regex_pattern <- RegexPattern(@String)

    # Boolean patterns use hard-coded "true" & "false" identifiers, even though
    # booleans literals are not built-in in Langkit but just a regular enum.
    bool_pattern <- or(
        | BoolPattern.True(@Identifier("true"))
        | BoolPattern.False(@Identifier("false"))
    )

    ellipsis_pattern <- EllipsisPattern(?pick(id "@") "...")
    integer_pattern <- IntegerPattern(@Number)
    list_pattern <- ListPattern(
        "[" list+(or(ellipsis_pattern | neg_pattern), ",") "]"
    )
    pattern_arg <- or(
        | FieldPatternDetail(id ":" / pattern)
        | PropertyPatternDetail(
            CallExpr(ref_id "(" / args ")") ":" / pattern
        )
    )
    expr <- or(
        | BinOp(stream_concat Op.StreamConcat(":::") expr)
        | stream_concat
    )
    stream_concat <- or(
        | BinOp(logic Op.StreamCons("::") stream_concat)
        | logic
    )
    logic <- or(
        | BinOp(
            logic
            or(
                | Op.OrInt("or" "?")
                | Op.Or("or")
                | Op.And("and")
                | Op.LogicAnd("%" "and")
                | Op.LogicOr("%" "or")
            )
            rel
        )
        | rel
    )
    rel <- or(NotExpr("not" eq) | eq)
    eq <- or(
        | BinOp(
            eq
            or(
                | Op.Lte("<=")
                | Op.Lt("<")
                | Op.Gte(">=")
                | Op.Gt(">")
                | Op.Eq("==")
                | Op.Ne("!=")
            )
            arith_1
        )
        | arith_1
    )
    arith_1 <- or(
        | BinOp(arith_1 or(Op.Plus("+") | Op.Minus("-") | Op.Amp("&")) arith_2)
        | arith_2
    )
    arith_2 <- or(
        | BinOp(arith_2 or(Op.Mult("*") | Op.Div("/")) arith_3)
        | arith_3
    )
    arith_3 <- or(
        | UnOp(or(Op.Plus("+") | Op.Minus("-")) isa_or_primary)
        | isa_or_primary
    )
    isa_or_primary <- or(
        | Isa(primary "is" / pattern)
        | AnyOf(primary "in" / AnyOfList+(primary, "|"))
        | LogicUnify(isa_or_primary "<->" / primary)
        | LogicPropagate(isa_or_primary "<-" logic_propagate_call)
        | LogicAssign(isa_or_primary "<-" / primary)
        | primary
    )
    logic_propagate_call <- LogicPropagateCall(callable_ref "%" "(" args ")")
    primary <- or(| lambda_expr | if_expr | raise_expr | try_expr | basic_expr)
    match_expr <- MatchExpr(
        "match"
        /
        expr
        "{"
        list+(or(
            | PatternMatchBranch("case" pattern "=>" / expr)
            | MatchBranch(
                "case" MatchValDecl(def_id ?pick(":" / type_ref)) "=>" / expr
            )
        ))
        "}"
    )
    num_lit <- NumLit(@Number)
    big_num_lit <- BigNumLit(@BigNumber)
    string_lit <- or(
        | SingleLineStringLit(@String)
        | PatternSingleLineStringLit(@PString)
        | block_string_lit
    )
    block_string_lit <- BlockStringLit(
        list+(BlockStringLine(@BlockStringLine))
    )
    char_lit <- CharLit(@Char)
    if_expr <- IfExpr(
        "if"
        /
        expr
        "then"
        expr
        list*(ElsifBranch("elif" / expr "then" expr))
        "else"
        expr
    )
    raise_expr <- RaiseExpr("raise" / ?pick("[" / type_ref "]") expr)
    try_expr <- TryExpr("try" / expr ?pick("else" / expr))
    array_literal <- ArrayLiteral(
        "[" / list*(expr, ",") "]" ?pick(":" / type_ref)
    )
    callable_ref <- or(
        DotExpr(callable_ref null(NullCondQualifier) "." ref_id) | ref_id
    )
    null_cond_qual <- NullCondQualifier("?")
    basic_expr <- or(
        | CallExpr(basic_expr "(" / args ")")
        | GenericInstantiation(basic_expr "[" type_list "]")
        | SubscriptExpr(basic_expr null_cond_qual "[" / expr "]")
        | ErrorOnNull(basic_expr "!")
        | KeepExpr(
            basic_expr null_cond_qual "." @Identifier("keep") "[" type_ref "]"
        )
        | CastExpr(
            basic_expr
            null_cond_qual
            "."
            @Identifier("as")
            /
            ExcludesNull("!")
            "["
            type_ref
            "]"
        )
        | LogicPredicate(basic_expr "%" "(" / args ")")
        | DotExpr(basic_expr null_cond_qual "." / ref_id)
        | LogicExpr("%" CallExpr(ref_id "(" / args ")"))
        | LogicExpr("%" / ref_id)
        | term
    )
    term <- or(
        | ParenExpr("(" / expr ")")
        | match_expr
        | null_lit
        | ref_id
        | block
        | num_lit
        | big_num_lit
        | string_lit
        | char_lit
        | query_comprehension
        | array_literal
    )
    basic_name <- or(
        DotExpr(basic_name null(NullCondQualifier) "." / ref_id) | ref_id
    )
    lambda_expr <- LambdaExpr(
        "(" lambda_param_list ")" ?pick(":" type_ref) "=>" / expr
    )
    null_lit <- NullLit("null" / ?pick("[" / type_ref "]"))
    argument <- Argument(?pick(ref_id "=") expr)
    args <- list*(argument, ",")
    decl_annotation_args <- ?DeclAnnotationArgs("(" args ")")
    decl_annotation <- DeclAnnotation("@" / id decl_annotation_args)
    query_comprehension <- Query(
        "from" expr
        "match" pattern
        ?pick("select" expr)
        ?pick("if" expr)
    )
}