from lark import Lark

GRAMMAR = """
FN: "fn"
IMPORT: "import"
COMMA: ","
DOT: "."
SINGLE_COMMENT: "//"
NEWLINE: /(\\r?\\n)+\\s*/
ANY: /.+/
WHITESPACE: " "
INTEGER: /[0-9]+/
ARROW: "->"
COM_START: "/*"
COM_END: "*/"
QUOTE: "\\""

identifier: WHITESPACE* ANY WHITESPACE*

single_comment: SINGLE_COMMENT ANY* NEWLINE
multi_comment: COM_START ANY* COM_END

import_stmt: IMPORT identifier NEWLINE

fn_arg: identifier identifier
parameters: fn_arg (COMMA fn_arg)
fn_stmt: FN identifier? "(" parameters? ")" [ARROW identifier] "{" NEWLINE? [stmt NEWLINE]* "}"

sign_int: "+" | "-"
string: QUOTE ANY* QUOTE
value: (sign_int* INTEGER) | string

call_stmt: [identifier DOT] identifier "(" [value COMMA]* ")"

stmt: value | import_stmt | fn_stmt | call_stmt

start: (NEWLINE | stmt)*
"""

def parse(string: str):
    """Parse using Lark"""
    parser = Lark(GRAMMAR, parser='lalr', debug=True)
    return parser.parse(string)