import random

from rply.token import BaseBox

def rep_join(l):
    format_string = ',\n'.join(
        [repr(i) if not isinstance(i, str) else i for i in l]
    ).replace('\n', '\n  ')

    if format_string != "":
        format_string = f"\n  {format_string}\n"

    return format_string

class Exp(BaseBox):
    def eval(self, vtable, ftable):
        return vtable, ftable, None

class ExpNumeral(Exp):
    def __init__(self, value: int):
        self.value = value

    def eval(self, vtable, ftable):
        return vtable, ftable, self.value

class ExpString(Exp):
    def __init__(self, value: str):
        self.value = value

    def eval(self, vtable, ftable):
        return vtable, ftable, self.value

class ExpVariable(Exp):
    def __init__(self, variable_id: str):
        self.variable_name = variable_id

    def eval(self, vtable, ftable):
        return vtable, ftable, vtable[self.variable_name]

class ExpABinop(Exp):
    def __init__(self,op: str,exp1: Exp, exp2: Exp):
        self.op = op
        self.exp1 = exp1
        self.exp2 = exp2

    def eval(self, vtable, ftable):
        if self.op == "+":
            vtable, ftable, r1 = self.exp1.eval(vtable, ftable)
            vtable, ftable, r2 = self.exp2.eval(vtable, ftable)
            return vtable, ftable, r1+r2
        elif self.op == "-":
            vtable, ftable, r1 = self.exp1.eval(vtable, ftable)
            vtable, ftable, r2 = self.exp2.eval(vtable, ftable)
            return vtable, ftable, r1-r2
        elif self.op == "*":
            vtable, ftable, r1 = self.exp1.eval(vtable, ftable)
            vtable, ftable, r2 = self.exp2.eval(vtable, ftable)
            return vtable, ftable, r1*r2
        elif self.op == "/":
            vtable, ftable, r1 = self.exp1.eval(vtable, ftable)
            vtable, ftable, r2 = self.exp2.eval(vtable, ftable)
            return vtable, ftable, r1/r2
        elif self.op == "=":
            vtable, ftable, r1 = self.exp1.eval(vtable, ftable)
            vtable, ftable, r2 = self.exp2.eval(vtable, ftable)
            return vtable, ftable, r1 == r2
        else:
            raise Exception(f"Binop {self.op} not implemented")

class ExpArg(Exp):
    def __init__(self, arg: int):
        self.arg = arg

    def eval(self, vtable, ftable):
        n = vtable["#ARGS"] - self.arg

        return vtable, ftable, vtable[f"#{n}"]

class Command(BaseBox):
    def eval(self, vtable, ftable):
        return vtable, ftable, None

class Builtin(Command):
    def __init__(self, builtin: str, args: list[Exp]):
        self.builtin = builtin
        self.args = args

    def eval(self, vtable, ftable):
        if self.builtin == "print":
            prints = []
            for arg in self.args:
                vtable, ftable, result = arg.eval(vtable, ftable)
                prints.append(str(result))

            print("   ".join(prints))

            return vtable, ftable, None
        else:
            raise Exception(f"Unknown builtin {self.builtin}")

class Do(BaseBox):
    def __init__(self, command: Command):
        self.command = command

    def eval(self, vtable, ftable):
        return self.command.eval(vtable, ftable)

class Statement(BaseBox):
    def eval(self, vtable, ftable):
        return vtable, ftable, None

class StatementSet(Statement):
    def __init__(self, expression: Exp, variable: str):
        self.expression = expression
        self.variable_name = variable

    def eval(self, vtable, ftable):
        vtable, ftable, result = self.expression.eval(
            vtable, ftable
        )
        vtable[self.variable_name] = result
        return vtable, ftable, None

class StatementDefine(Statement):
    def __init__(self, function_name : str, parameters: int, statement: Statement):
        self.function_name = function_name
        self.statement = statement
        self.parameters = parameters

    def eval(self, vtable, ftable):
        ftable[self.function_name] = (self.parameters, self.statement)
        return vtable, ftable, None

class StatementReturn(Statement):
    def __init__(self, expression: Exp):
        self.value = expression

    def eval(self, vtable, ftable):
        vtable, ftable, result = self.value.eval(vtable, ftable)
        vtable["#return"] = result
        return vtable, ftable, result

class Scope(Statement):
    def __init__(self, statements: list[Statement]):
        self.statements = statements

    def eval(self, vtable, ftable):
        result = None
        for statement in self.statements:
            vtable, ftable, _ = statement.eval(vtable, ftable)
            if "#return" in vtable:
                result = vtable["#return"]
                del vtable["#return"]
                break

        return vtable, ftable, result

class Call(Statement):
    def __init__(self, function_name: str, arguments: list[Exp]):
        self.function_name = function_name
        self.arguments = arguments

    def eval(self, vtable, ftable):
        assert len(self.arguments) == ftable[self.function_name][0]

        args = []
        for argument in self.arguments:
            vtable, ftable, result = argument.eval(vtable,ftable)
            args.append(result)

        n = vtable["#ARGS"]-1

        for i, arg in enumerate(args):
            vtable[f"#{n+len(args)-i}"] = arg

        vtable["#ARGS"] += len(args)

        vtable, ftable, result = ftable[self.function_name][1].eval(vtable, ftable)

        if "#return" in vtable:
            del vtable["#return"]

        vtable["#ARGS"] -= len(args)

        return vtable, ftable, result

class StatementIf(Statement):
    def __init__(self, statement: Statement, condition):
        self.statement = statement
        self.condition = condition

    def eval(self, vtable, ftable):
        vtable, ftable, result = self.condition.eval(vtable, ftable)
        if result:
            return self.statement.eval(vtable, ftable)
        else:
            return vtable, ftable, None

class Maybe(BaseBox):
    def __init__(self, statement : Statement):
        self.statement = statement

    def eval(self, vtable, ftable):
        if random.randint(0,1) == 1:
            return self.statement.eval(vtable,ftable)
        else:
            return vtable, ftable, None

class Program(BaseBox):
    def __init__(self, statements: list[Statement]) -> None:
        self.statements = statements

    def __repr__(self) -> str:
        statements_string = f"statements([{rep_join(self.statements)}])"
        return statements_string

    def eval(self, *_):
        vtable = {"#ARGS": 0}
        ftable = {}

        for statement in self.statements:
            vtable, ftable, _ = statement.eval(vtable, ftable)
            if "#return" in vtable:
                break