from pyparsing import *

ParserElement.enable_packrat()

COMMENT = Regex(r"/\-(?:[^-]|\-(?!/))*\-\/").set_name("comment")

IDENT = unicode_set.identifier()

DEF, EXAMPLE, IND, WHERE, OPEN, TYPE, NOMATCH, MATCH, WITH, UNDER, CLASS, INST = map(
    lambda w: Suppress(Keyword(w)),
    "def example inductive where open Type nomatch match with _ class instance".split(),
)

ASSIGN, ARROW, FUN, TO = map(lambda s: Suppress(s[0]) | Suppress(s[1:]), "≔:= →-> λfun ↦=>".split())

LPAREN, RPAREN, LBRACE, RBRACE, LBRACKET, RBRACKET, COLON, BAR, NEWLINE = map(Suppress, "(){}[]:|\n")
INLINE_WHITE = Opt(Suppress(White(" \t\r"))).set_name("inline_whitespace")

forwards = lambda names: map(lambda n: Forward().set_name(n), names.split())

expr, atom, fn_type, fn, match, nomatch, call, p_expr, type_, ph, ref = forwards(
    "expr atom fn_type fn match nomatch call paren_expr type placeholder ref"
)
case, i_arg, e_arg = forwards("case implicit_arg explicit_arg")

infix_op = lambda s: (one_of(s), 2, OpAssoc.LEFT)
expr <<= infix_notation(atom, [infix_op("* /"), infix_op("+ -")])
atom <<= fn_type | fn | match | nomatch | call | p_expr | type_ | ph | ref

name = Group(IDENT).set_name("name")
i_param = (LBRACE + name + COLON + expr + RBRACE).set_name("implicit_param")
e_param = (LPAREN + name + COLON + expr + RPAREN).set_name("explicit_param")
c_param = (LBRACKET + name + COLON + expr + RBRACKET).set_name("class_param")
param = (i_param | e_param | c_param).set_name("param")
fn_type <<= param + ARROW + expr
fn <<= FUN - Group(OneOrMore(name)) + TO + expr
match <<= MATCH - (type_ | ref | p_expr) + WITH + Group(OneOrMore(case))
case <<= BAR - ref + Group(ZeroOrMore(name)) + TO + expr
nomatch <<= (NOMATCH - INLINE_WHITE + e_arg).leave_whitespace()
callee = ref | p_expr
call <<= (callee + OneOrMore(INLINE_WHITE + (i_arg | e_arg))).leave_whitespace()
i_arg <<= LPAREN + IDENT + ASSIGN + expr + RPAREN
e_arg <<= (type_ | ph | ref | p_expr).leave_whitespace()
p_expr <<= LPAREN + expr + RPAREN
type_ <<= Group(TYPE)
ph <<= Group(UNDER)
ref <<= Group(name)

return_type = Opt(COLON + expr)
params = Group(ZeroOrMore(param))
def_ = (DEF - ref + params + return_type + ASSIGN + expr).set_name("definition")
example = (EXAMPLE - params + return_type + ASSIGN + expr).set_name("example")
type_arg = (LPAREN + ref + ASSIGN + expr + RPAREN).set_name("type_arg")
ctor = (BAR - ref + params + Group(ZeroOrMore(type_arg))).set_name("constructor")
data = (IND - ref + params + WHERE + Group(ZeroOrMore(ctor)) + OPEN + IDENT).set_name("datatype")
c_field = (name + COLON + expr).set_name("class_field")
class_ = (CLASS - ref + params + WHERE + Group(ZeroOrMore(c_field)) + OPEN + IDENT).set_name("class")
i_field = (ref + ASSIGN + expr).set_name("instance_field")
inst = (INST - COLON + expr + WHERE + Group(ZeroOrMore(i_field))).set_name("instance")
declaration = (def_ | example | data | class_ | inst).set_name("declaration")

program = ZeroOrMore(declaration).ignore(COMMENT).set_name("program")

line_exact = lambda w: Suppress(AtLineStart(w) + LineEnd())
markdown = line_exact("```lean") + program + line_exact("```")
