from graphviz import Digraph
from lark import Lark, Tree, Token
from lark import Transformer, v_args
from lark.visitors import Interpreter
import os


# 定义语法规则
grammar = """
?start: (stmt _NL?)*

?stmt: PRINT expr                          -> print_stmt
     | NAME EQUAL expr                     -> assign_stmt
     | NAME PLUS_EQUAL expr                -> add_assign_stmt
     | NAME MINUS_EQUAL expr               -> subtract_assign_stmt
     | IF expr THEN stmt+ ELSE stmt+ END   -> if_stmt
     | DO stmt+ END WHILE expr             -> do_while_stmt  # 确保 'do-while' 循环有 'END'
     | expr                                -> expr_stmt
     | _NL                                 -> newline  // 新增空行处理

?expr: or_expr

?or_expr: and_expr
        | or_expr "or" and_expr            -> logical_or

?and_expr: equality_expr
         | and_expr "and" equality_expr    -> logical_and

?equality_expr: comparison_expr
              | equality_expr EQ comparison_expr  -> eq
              | equality_expr NE comparison_expr  -> ne

?comparison_expr: term
                | comparison_expr GT term                  -> gt
                | comparison_expr GE term                  -> ge
                | comparison_expr LE term                  -> le
                | comparison_expr LT term                  -> lt

?term: factor
     | term "*" factor               -> mul
     | term "/" factor               -> div
     | term "%" factor               -> mod
     | term "**" factor              -> pow

?factor: NUMBER                      
       | NAME                        -> var
       | "(" expr ")"                -> grouped_expr

%import common.CNAME -> NAME
%import common.NUMBER
%import common.WS
%import common.NEWLINE -> _NL  // 直接使用 common.NEWLINE 并重命名为 _NL

// 忽略空白和注释
%ignore WS
%ignore /#.*/

PRINT: "print"
EQUAL: "="
PLUS_EQUAL: "+="
MINUS_EQUAL: "-="
IF: "if"
THEN: "then"
ELSE: "else"
END: "end"
DO: "do"
WHILE: "while"
GT: ">"
GE: ">="
EQ: "=="
NE: "!="
LE: "<="
LT: "<"
MOD: "%"
POW: "**"
"""


# 定义解释器类
class SimpleLangInterpreter(Interpreter):
    def __init__(self):
        super().__init__()
        self.variables = {}

    def visit(self, tree_or_token):
        if isinstance(tree_or_token, Tree):
            return super().visit(tree_or_token)
        elif isinstance(tree_or_token, Token):
            return self.visit_token(tree_or_token)
        else:
            raise ValueError(f"Unexpected type: {type(tree_or_token)}")

    def visit_token(self, token):
        if token.type == "NUMBER":
            return int(token.value)
        elif token.type == "NAME":
            name = str(token.value)
            if name in self.variables:
                return self.variables[name]
            else:
                raise ValueError(f"Undefined variable '{name}'")
        elif token.type == "BOOLEAN":
            return token.value == "True"
        else:
            raise ValueError(f"Unrecognized token type: {token.type}")

    def assign_stmt(self, tree):
        name = str(tree.children[0])
        value = self.visit(tree.children[2])
        self.variables[name] = value

    def add_assign_stmt(self, tree):
        name = str(tree.children[0])
        if name in self.variables:
            current_value = self.variables[name]
            right_value = self.visit(tree.children[2])
            new_value = current_value + right_value
            self.variables[name] = new_value
        else:
            raise ValueError(f"Undefined variable '{name}'")

    def subtract_assign_stmt(self, tree):
        name = str(tree.children[0])
        if name in self.variables:
            current_value = self.variables[name]
            right_value = self.visit(tree.children[2])
            new_value = current_value - right_value
            self.variables[name] = new_value
        else:
            raise ValueError(f"Undefined variable '{name}'")

    def print_stmt(self, tree):
        value = self.visit(tree.children[1])
        print(value)

    def if_stmt(self, tree):
        condition = self.visit(tree.children[1])
        if condition:
            self.visit_children(tree.children[3])
        else:
            self.visit_children(tree.children[5])

    def do_while_stmt(self, tree):
        while True:
            for child in tree.children[1:-1]:
                if isinstance(child, Tree):
                    self.visit(child)
            condition = self.visit(tree.children[-1])
            if not condition:
                break

    def logical_or(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left or right

    def logical_and(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left and right

    def eq(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left == right

    def ne(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left!= right

    def le(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left <= right

    def lt(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left < right

    def ge(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left >= right

    def gt(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[2])
        return left > right

    def mul(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left * right

    def div(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left / right

    def mod(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left % right

    def pow(self, tree):
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left ** right

    def var(self, tree):
        name = str(tree.children[0])
        if name in self.variables:
            return self.variables[name]
        else:
            raise ValueError(f"Undefined variable '{name}'")

    def expr_stmt(self, tree):
        result = self.visit(tree.children[0])
        return result


# 解析代码
def parse_code(code, grammar):
    try:
        parser = Lark(grammar, parser='lalr', lexer='basic')
        return parser.parse(code)
    except Exception as e:
        print(f"解析错误: {e}")
        return None


# 可视化解析树
def draw_tree(tree, graph, parent=None, count=0):
    node_id = str(count)
    label = tree.data if isinstance(tree, Tree) else str(tree)
    graph.node(node_id, label)

    if parent is not None:
        graph.edge(parent, node_id)

    count += 1
    for child in tree.children:
        if isinstance(child, Tree):
            count = draw_tree(child, graph, node_id, count)
        elif isinstance(child, Token):
            leaf_id = str(count)
            graph.node(leaf_id, str(child))
            graph.edge(node_id, leaf_id)
            count += 1
    return count


def visualize_tree(tree):
    if tree:
        graph = Digraph(format="png")
        draw_tree(tree, graph)
        return graph
    return None


# 解释并执行代码
def interpret_code(tree):
    if tree:
        interpreter = SimpleLangInterpreter()
        try:
            interpreter.visit(tree)
        except Exception as e:
            print(f"解释执行错误: {e}")


if __name__ == "__main__":
    code = """
    x = 10
    y = 5
    z = 2

    print x / y 
    print x ** y  
    print x % y  

    if x <= y then
     print x
    else
     print y
    end

    i = 0
    do
      print i
      i += 1
    end  
    while i < 5 
    """

    parsed_tree = parse_code(code, grammar)
    tree_graph = visualize_tree(parsed_tree)
    if tree_graph:
        os.environ["PATH"] += os.pathsep + "/usr/local/bin"
        tree_graph.render('simple_lang_tree_demo_exp1')
        tree_graph.view()

    interpret_code(parsed_tree)

这段代码实现了一个简单编程语言的解析、可视化解析树以及解释执行的功能。具体如下：

1. 定义语法规则

使用 Lark 库定义了一种简单编程语言的语法规则。该语法支持以下语句和表达式：

• 语句：
  • 打印语句 (PRINT expr)
  • 赋值语句 (NAME EQUAL expr)
  • 加法赋值语句 (NAME PLUS_EQUAL expr)
  • 减法赋值语句 (NAME MINUS_EQUAL expr)
  • if 语句 (IF expr THEN stmt+ ELSE stmt+ END)
  • do-while 循环语句 (DO stmt+ END WHILE expr)
  • 表达式语句 (expr)
  • 空行 (_NL)
• 表达式：
  • 逻辑或表达式 (or_expr "or" and_expr)
  • 逻辑与表达式 (and_expr "and" equality_expr)
  • 相等性表达式 (equality_expr EQ comparison_expr 或 equality_expr NE comparison_expr)
  • 比较表达式 (comparison_expr GT term 等各种比较操作)
  • 项 (term "*" factor 等各种算术操作)
  • 因子 (NUMBER、NAME 或 ("(" expr ")"))

2. 定义解释器类

SimpleLangInterpreter 类继承自 Interpreter，实现了对上述语法规则的解释执行。它维护了一个变量字典 self.variables 来存储变量的值，并针对每种语句和表达式类型定义了相应的处理方法，例如：

• assign_stmt：处理赋值语句，将表达式的值赋给变量。
• add_assign_stmt 和 subtract_assign_stmt：处理加法和减法赋值语句。
• print_stmt：处理打印语句，输出表达式的值。
• if_stmt：处理 if 语句，根据条件执行相应的语句块。
• do_while_stmt：处理 do-while 循环语句。
• 各种表达式操作方法，如 logical_or、logical_and、eq、ne 等，用于计算表达式的值。

3. 解析代码

parse_code 函数使用 Lark 解析器对输入的代码进行解析。如果解析成功，返回解析树；如果解析失败，打印错误信息。

4. 可视化解析树

draw_tree 函数递归地遍历解析树，将每个节点和边添加到 graphviz 的 Digraph 对象中，从而实现解析树的可视化。visualize_tree 函数调用 draw_tree 并返回生成的 Digraph 对象。

5. 解释并执行代码

interpret_code 函数创建一个 SimpleLangInterpreter 实例，并使用它来访问解析树，从而解释执行代码。如果在解释执行过程中发生错误，打印错误信息。

6. 主程序

在 if __name__ == "__main__": 块中，定义了一段示例代码，然后：

• 调用 parse_code 解析代码，得到解析树。
• 调用 visualize_tree 可视化解析树，并将生成的图形保存为 simple_lang_tree_demo_exp1，然后打开查看。
• 调用 interpret_code 解释执行代码。

总体而言，这段代码提供了一个简单的编译器前端，包括语法分析和语义解释，并且能够可视化解析过程的中间结果——解析树。