from jf.parser.lexer import TokenType
from jf.parser.parser_base import *
from jf.ast.toplevel import *
from jf.ast.common import *
from jf.ast.expr import *
from jf.ast.stmt import *
from jf.utils import *
import functools
from enum import IntEnum


class PrecLevel(IntEnum):
    ALL = 0

    LogicOr = 5
    LogicAnd = 10
    Comparison = 20
    Multiplication = 30
    Addition = 40

    Prefix = 60

    Indexing = 80

    Call = 90


@mparser
def block(parser):
    yield exact(TokenType.Op_LBrace)
    result = yield parser
    yield exact(TokenType.Op_RBrace)
    return result


def block_or_semicolon(contents):
    return choice(
        exact(TokenType.Op_Semicolon).bind0(value(None)),
        block(contents)
    )


def top_level_def():
    return choice(
        import_decl(),
        namespace(),
        struct(),
        alias(),
        function()
    )


def struct_def():
    return choice(
        alias(),
        field(),
        function(),
        case()
    )


def struct_defs():
    return many(struct_def())


def top_level_defs():
    return many(top_level_def())


@mparser
def namespace():
    loc = yield get_loc()
    yield exact(TokenType.K_Namespace)
    name = yield exact(TokenType.Identifier)
    contents = yield block_or_semicolon(top_level_defs())
    return Namespace(loc, name, contents)


@functools.lru_cache(maxsize=None)
def base_expression(level):
    prefix_ops = []
    for token, op in [
        (TokenType.K_Ref, UnOp.RefRef),
        (TokenType.K_Ptr, UnOp.RefPtr),
        (TokenType.K_RefIn, UnOp.RefIn),
        (TokenType.K_RefOut, UnOp.RefOut),
        (TokenType.Op_Star, UnOp.Deref),
        (TokenType.K_Not, UnOp.LogicNot),
        (TokenType.Op_Minus, UnOp.Negate),
    ]:
        (lambda op:
            prefix_ops.append(
                exact(token)
                .bind0(expression(max(level, PrecLevel.Prefix)))
                .bind(lambda ex: value(ExprUnOp(ex, op)))
            )
        )(op)

    decode_ref_type = choice(
        exact(TokenType.K_Ref).bind0(value(UnOp.AddressRef)),
        exact(TokenType.K_Ptr).bind0(value(UnOp.AddressPtr)),
        exact(TokenType.K_RefIn).bind0(value(UnOp.AddressRefIn)),
        exact(TokenType.K_RefOut).bind0(value(UnOp.AddressRefOut)),
        value(UnOp.AddressRef)
    )

    prefix_ops.append(
        exact(TokenType.Op_Amp)
        .bind0(decode_ref_type)
        .bind(lambda op:
            expression(max(level, PrecLevel.Prefix))
            .bind(lambda ex: value(ExprUnOp(ex, op)))
        )
    )

    prefix_ops.append(
        between(
            exact(TokenType.Op_LBracket),
            exact(TokenType.Op_RBracket),
            optional(expression())
        )
        .bind(lambda size_ex:
            expression(max(level, PrecLevel.Prefix))
            .bind(lambda ex: value(ExprArray(ex, size_ex)))
        )
    )

    return choice(
        between(exact(TokenType.Op_LParen), exact(TokenType.Op_RParen), expression()),
        exact(TokenType.Identifier).bind(lambda name: value(ExprID(name))),
        exact(TokenType.FloatLiteral).bind(lambda content: value(ExprNum(content))),
        exact(TokenType.IntLiteral).bind(lambda content: value(ExprNum(content))),
        exact(TokenType.StringLiteral).bind(lambda content: value(ExprString(content))),
        *prefix_ops
    )


@functools.lru_cache(maxsize=None)
def expression_rhs(level):
    parsers_list = []

    if level <= PrecLevel.Indexing:
        parsers_list.append(
            exact(TokenType.Op_Dot)
            .bind0(exact(TokenType.Identifier))
            .bind(lambda id: value(lambda lhs: ExprField(lhs, id)))
        )
    if level <= PrecLevel.LogicOr:
        parsers_list.append(
            exact(TokenType.K_Or)
            .bind0(expression(PrecLevel.LogicOr + 1))
            .bind(lambda rhs: value(lambda lhs: ExprBinOp(lhs, BinOp.LogicOr, rhs)))
        )
    if level <= PrecLevel.LogicAnd:
        parsers_list.append(
            exact(TokenType.K_And)
            .bind0(expression(PrecLevel.LogicAnd + 1))
            .bind(lambda rhs: value(lambda lhs: ExprBinOp(lhs, BinOp.LogicAnd, rhs)))
        )
    if level <= PrecLevel.Comparison:
        for token, op in [
            (TokenType.Op_EQ, BinOp.EQ),
            (TokenType.Op_NEQ, BinOp.NEQ),
            (TokenType.Op_LE, BinOp.LE),
            (TokenType.Op_LT, BinOp.LT),
            (TokenType.Op_GE, BinOp.GE),
            (TokenType.Op_GT, BinOp.GT),
        ]:
            (lambda op:
            parsers_list.append(
                exact(token)
                .bind0(expression(PrecLevel.Comparison + 1))
                .bind(lambda rhs: value(lambda lhs: ExprBinOp(lhs, op, rhs)))
            ))(op)
    if level <= PrecLevel.Multiplication:
        for token, op in [
            (TokenType.Op_Star, BinOp.Multiply),
            (TokenType.Op_Div, BinOp.Divide),
        ]:
            (lambda op:
            parsers_list.append(
                exact(token)
                .bind0(expression(PrecLevel.Multiplication + 1))
                .bind(lambda rhs: value(lambda lhs: ExprBinOp(lhs, op, rhs)))
            ))(op)
    if level <= PrecLevel.Addition:
        for token, op in [
            (TokenType.Op_Plus, BinOp.Add),
            (TokenType.Op_Minus, BinOp.Subtract),
        ]:
            (lambda op:
            parsers_list.append(
                exact(token)
                .bind0(expression(PrecLevel.Addition + 1))
                .bind(lambda rhs: value(lambda lhs: ExprBinOp(lhs, op, rhs)))
            ))(op)

    if level <= PrecLevel.Call:
        parsers_list.append(
            between(
                exact(TokenType.Op_LParen),
                exact(TokenType.Op_RParen),
                sep_by(exact(TokenType.Op_Comma), expression())
            )
            .bind(lambda args: value(lambda lhs: ExprCall(lhs, args)))
        )
        parsers_list.append(
            between(
                exact(TokenType.Op_LBracket),
                exact(TokenType.Op_RBracket),
                sep_by(exact(TokenType.Op_Comma), expression())
            )
            .bind(lambda args: value(lambda lhs: ExprSubscript(lhs, args)))
        )

    def fn(expr):
        return choice(
            choice(
                *parsers_list
            ).bind(lambda f: fn(f(expr))),
            value(expr)
        )
    return fn


@mparser
def expression(level=PrecLevel.ALL):
    prim = yield base_expression(level)

    return expression_rhs(level)(prim)


@mparser
def type_binder():
    yield exact(TokenType.Op_Colon)
    type = yield expression()
    return type


def maybe_type_binder():
    return optional(type_binder())


@mparser
def param_group():
    names = yield many1(exact(TokenType.Identifier))
    type = yield maybe_type_binder()
    return ParamGroup(names, type)


@mparser
def param_list():
    groups = yield sep_by(
        exact(TokenType.Op_Comma),
        param_group()
    )
    return ParamList(groups)


@mparser
def generic_params():
    yield exact(TokenType.Op_LBracket)
    params = yield param_list()
    yield exact(TokenType.Op_RBracket)
    return params


def maybe_generic_params():
    return optional(generic_params())


@mparser
def function_params():
    yield exact(TokenType.Op_LParen)
    params = yield param_list()
    yield exact(TokenType.Op_RParen)
    return params


@mparser
def constraint():
    return (yield expression())


@mparser
def where_clause():
    yield exact(TokenType.K_Where)
    constraints = sep_by(exact(TokenType.Op_Comma), constraint())
    return constraints


def maybe_where_clause():
    return optional(where_clause())


@mparser
def return_statement():
    loc = yield get_loc()
    yield exact(TokenType.K_Return)
    expr = yield expression()
    yield exact(TokenType.Op_Semicolon)
    return StmtReturn(loc, expr)


@mparser
def expression_statement():
    loc = yield get_loc()
    expr = yield expression()
    result, assign = yield optional_with_result(assign_statement(loc, expr))
    if result:
        return assign
    yield exact(TokenType.Op_Semicolon)
    return StmtExpr(loc, expr)


@mparser
def assign_statement(loc, lhs):
    yield exact(TokenType.Op_Assign)
    rhs = yield expression()
    yield exact(TokenType.Op_Semicolon)
    return StmtAssign(loc, lhs, rhs)


@mparser
def while_statement():
    loc = yield get_loc()
    yield exact(TokenType.K_While)
    cond = yield expression()
    stmts = yield block(statements())
    return StmtWhile(loc, cond, stmts)


@mparser
def if_branch():
    loc = yield get_loc()
    yield exact(TokenType.K_If)
    cond = yield expression()
    stmts = yield block(statements())
    return IfBranch(loc, cond, stmts)


@mparser
def if_statement():
    loc = yield get_loc()
    branches = [(yield if_branch())]
    else_block = None
    while True:
        if (yield optional(exact(TokenType.K_Else))):
            branch = yield optional(if_branch())
            if branch:
                branches.append(branch)
                continue
            else:
                else_block = yield block(statements())
                break
        else:
            break
    return StmtIf(loc, branches, else_block)


def maybe_initializer():
    return optional(exact(TokenType.Op_Assign).bind0(expression()))


@mparser
def var_statement():
    loc = yield get_loc()
    yield exact(TokenType.K_Var)
    names = yield many1(exact(TokenType.Identifier))
    type = yield maybe_type_binder()
    initializer = yield maybe_initializer()
    yield exact(TokenType.Op_Semicolon)
    return StmtVar(loc, names, type, initializer)


def statement():
    return choice(
        return_statement(),
        while_statement(),
        var_statement(),
        if_statement(),
        expression_statement()
    ).named("statement")


def statements():
    return many(statement())


def qual_id():
    return sep_by1(exact(TokenType.Op_Dot), exact(TokenType.Identifier))


def maybe_qual_id():
    return optional(qual_id())


@mparser
def function_body():
    stmts = yield block(statements())
    return FunctionBody(stmts)


@mparser
def function_ffi():
    loc = yield get_loc()
    yield exact(TokenType.Op_Assign)
    yield exact(TokenType.K_Import)
    string = yield exact(TokenType.StringLiteral)
    yield exact(TokenType.Op_Semicolon)
    return FunctionFFI(loc, string)


def function_def():
    return choice(
        exact(TokenType.Op_Semicolon).bind0(value(None)),
        function_body(),
        function_ffi()
    )


@mparser
def field():
    loc = yield get_loc()
    yield exact(TokenType.K_Field)
    names = yield many1(exact(TokenType.Identifier))
    type = yield type_binder()
    yield exact(TokenType.Op_Semicolon)
    return StructField(loc, names, type)


@mparser
def case():
    loc = yield get_loc()
    yield exact(TokenType.K_Case)
    name = yield exact(TokenType.Identifier)
    params = yield optional(function_params())
    type = yield maybe_type_binder()
    defs = yield block_or_semicolon(struct_defs())
    return StructCase(loc, name, params, type, defs)


@mparser
def struct():
    loc = yield get_loc()
    yield exact(TokenType.K_Struct)
    name = yield exact(TokenType.Identifier)
    params = yield maybe_generic_params()
    kind = yield maybe_type_binder()
    where = yield maybe_where_clause()
    defs = yield block_or_semicolon(struct_defs())
    return Struct(loc, name, params, kind, where, defs)


@mparser
def alias():
    loc = yield get_loc()
    yield exact(TokenType.K_Alias)
    name = yield exact(TokenType.Identifier)
    params = yield maybe_generic_params()
    yield exact(TokenType.Op_Assign)
    expr = yield expression()
    yield exact(TokenType.Op_Semicolon)
    return Alias(loc, name, params, expr)


@mparser
def function():
    loc = yield get_loc()
    kind = yield choice(exact(TokenType.K_Function), exact(TokenType.K_Method))
    name = yield exact(TokenType.Identifier)
    gen_params = yield maybe_generic_params()
    params = yield function_params()
    ret_ty = yield maybe_type_binder()
    where = yield maybe_where_clause()
    definition = yield function_def()
    return Function(loc, kind, name, gen_params, params, ret_ty, where, definition)


@mparser
def import_decl():
    loc = yield get_loc()
    yield exact(TokenType.K_Import)
    qid = yield maybe_qual_id()
    from_qid = yield optional(exact(TokenType.K_From).bind0(qual_id()))
    into_id = yield optional(exact(TokenType.K_Into).bind0(exact(TokenType.Identifier)))
    yield exact(TokenType.Op_Semicolon)
    return Import(loc, qid, from_qid, into_id)


@mparser
def file():
    defs = yield top_level_defs()
    yield exact(TokenType.EOF)
    return File(defs)