from stdint import *
import t, c
from t import CInt, CExport
from stdio import printf
import memhub
import stdlib
import sys
import string
import ast
import w32.win32console


# MBuddy arena 大小：1MB
ARENA_SIZE: t.CDefine = 1048576
# mpool 内存大小：512KB
POOL_SIZE: t.CDefine = 524288
# dump 缓冲区大小：256KB
DUMP_SIZE: t.CDefine = 262144


def classify_list(lst: list[ast.AST | t.CPtr] | t.CPtr, depth: t.CInt):
    """递归遍历 list[AST | t.CPtr] 容器，逐元素分类（orelse/args/comparators 等）"""
    if lst is None:
        return
    n: t.CSizeT = lst.__len__()
    i: t.CSizeT = 0
    while i < n:
        classify_print(lst.get(i), depth + 1)
        i += 1


def classify_print(node: ast.AST | t.CPtr, depth: t.CInt):
    """递归遍历 parse 树，用 isinstance 运行时判定每个节点的类别并打印。

    依赖 TPV isinstance 的 vtable 子类集合比较：
    - isinstance(node, ast.Stmt)     → 语句节点（Assign/If/For/While/...）
    - isinstance(node, ast.ExprNode) → 表达式节点（Name/Constant/BinOp/...）
    - 两者皆非                       → 其他（Module/Expression/...）
    验证运行时（非编译期静态）正确区分 Stmt 与 ExprNode。
    """
    if node is None:
        return
    sp: t.CInt = 0
    while sp < depth:
        printf("  ")
        sp += 1
    if isinstance(node, ast.Stmt) != 0:
        printf("STMT  %s\n", node.type_name())
    elif isinstance(node, ast.ExprNode) != 0:
        printf("EXPR  %s\n", node.type_name())
    else:
        printf("OTHER %s\n", node.type_name())
    # 按 kind 遍历语句持有的表达式字段（语句 body 块均在 children）。
    # 属性访问前先 cast 成具体节点类：node 静态类型是 ast.AST，
    # 直接访问 node.value 会在 AST 基类字段表 fast-fail（按需过滤原则）。
    k: t.CInt = node.kind()
    if k == ast.ASTKind.Assign:
        asg: ast.Assign | t.CPtr = (ast.Assign | t.CPtr)(node)
        classify_print(asg.value, depth + 1)
    elif k == ast.ASTKind.If:
        ifn: ast.If | t.CPtr = (ast.If | t.CPtr)(node)
        classify_print(ifn.test, depth + 1)
        classify_list(ifn.orelse, depth + 1)
    elif k == ast.ASTKind.While:
        whl: ast.While | t.CPtr = (ast.While | t.CPtr)(node)
        classify_print(whl.test, depth + 1)
        classify_list(whl.orelse, depth + 1)
    elif k == ast.ASTKind.For:
        for_nd: ast.For | t.CPtr = (ast.For | t.CPtr)(node)
        classify_print(for_nd.iter, depth + 1)
        classify_print(for_nd.target, depth + 1)
        classify_list(for_nd.orelse, depth + 1)
    elif k == ast.ASTKind.Return:
        ret_nd: ast.Return | t.CPtr = (ast.Return | t.CPtr)(node)
        classify_print(ret_nd.value, depth + 1)
    elif k == ast.ASTKind.Expr:
        expr_nd: ast.Expr | t.CPtr = (ast.Expr | t.CPtr)(node)
        classify_print(expr_nd.value, depth + 1)
    elif k == ast.ASTKind.BinOp:
        bin_nd: ast.BinOp | t.CPtr = (ast.BinOp | t.CPtr)(node)
        classify_print(bin_nd.left, depth + 1)
        classify_print(bin_nd.right, depth + 1)
    elif k == ast.ASTKind.UnaryOp:
        una_nd: ast.UnaryOp | t.CPtr = (ast.UnaryOp | t.CPtr)(node)
        classify_print(una_nd.operand, depth + 1)
    elif k == ast.ASTKind.Compare:
        cmp_nd: ast.Compare | t.CPtr = (ast.Compare | t.CPtr)(node)
        classify_print(cmp_nd.left, depth + 1)
        classify_list(cmp_nd.comparators, depth + 1)
    elif k == ast.ASTKind.BoolOp:
        bool_nd: ast.BoolOp | t.CPtr = (ast.BoolOp | t.CPtr)(node)
        classify_list(bool_nd.values, depth + 1)
    elif k == ast.ASTKind.Call:
        call_nd: ast.Call | t.CPtr = (ast.Call | t.CPtr)(node)
        classify_print(call_nd.func, depth + 1)
        classify_list(call_nd.args, depth + 1)
    else:
        pass
    # children（语句 body 块）
    if node.children is not None:
        n: t.CSizeT = node.children.__len__()
        i: t.CSizeT = 0
        while i < n:
            classify_print(node.children.get(i), depth + 1)
            i += 1


def main() -> CInt | CExport:
    w32.win32console.SetConsoleCP(65001)
    w32.win32console.SetConsoleOutputCP(65001)

    printf("AST Comprehensive Test Start\n")

    # 初始化全局 MBuddy 内存分配器
    arena: bytes = stdlib.malloc(ARENA_SIZE)
    if arena == None:
        printf("FAIL: arena alloc failed\n")
        return 1
    mbuddy_obj: memhub.MemBuddy | t.CPtr = memhub.MemBuddy(arena, ARENA_SIZE)
    if mbuddy_obj == None:
        printf("FAIL: mbuddy init failed\n")
        return 1

    sys._mbuddy = mbuddy_obj
    ast._mbuddy = mbuddy_obj

    # 为 ast 解析器创建 mpool
    pool_mem: bytes = stdlib.malloc(POOL_SIZE)
    if pool_mem == None:
        printf("FAIL: malloc for pool failed\n")
        return 1
    pool: memhub.MemPool | t.CPtr = memhub.MemPool(pool_mem, POOL_SIZE)
    if pool == None:
        printf("FAIL: memhub.MemPool init failed\n")
        return 1

    # dump 缓冲区（独立分配）
    dump_mem: bytes = stdlib.malloc(DUMP_SIZE)
    if dump_mem == None:
        printf("FAIL: malloc for dump failed\n")
        return 1
    dump_buf: t.CChar | t.CPtr = t.CChar(dump_mem, t.CPtr)

    # 综合测试源码：覆盖 CPython 主要语法
    src: str = """x = 1
y = 2
z = x + y
a = x - y
b = x * y
c = x / y
d = x == y
e = x < y
f = x > y
g = x and y
h = x or y
i = not x
def foo(a, b):
    return a + b
result = foo(x, y)
if x > 0:
    print(x)
else:
    print(y)
for i in range(10):
    print(i)
while x > 0:
    x = x - 1
class Point:
    pass
"""

    printf("Parsing comprehensive source...\n")

    # 解析
    ast._init_tables(pool)
    lx_dbg: t.CPtr = ast.new_lexer(pool)
    if lx_dbg == None:
        printf("FAIL: new_lexer returned None\n")
        stdlib.free(pool_mem)
        stdlib.free(dump_mem)
        return 1
    ast._lexer_init(lx_dbg, src, pool)
    tokens_dbg: t.CPtr = ast.tokenize(lx_dbg)
    if tokens_dbg == None:
        printf("FAIL: tokenize returned None\n")
        stdlib.free(pool_mem)
        stdlib.free(dump_mem)
        return 1
    tree: ast.AST | t.CPtr = ast.parse_tokens(pool, tokens_dbg)
    if tree == None:
        printf("FAIL: parse returned None\n")
        stdlib.free(pool_mem)
        stdlib.free(dump_mem)
        return 1

    printf("PASS: parse succeeded\n")

    # dump AST
    printf("Calling ast.dump...\n")
    ast.dump(tree, dump_buf, DUMP_SIZE)
    printf("ast.dump returned, now printing buffer...\n")
    printf("=== AST DUMP ===\n%s\n=== END ===\n", dump_buf)

    # ============================================================
    # Stmt/Expr 运行时区分测试（isinstance vtable 子类集合比较）
    # ============================================================
    printf("=== STMT/EXPR RUNTIME DISTINGUISH TEST ===\n")
    classify_print(tree, 0)

    # 断言 1: Module 既非 Stmt 亦非 ExprNode（顶层容器）
    if isinstance(tree, ast.Stmt) != 0 or isinstance(tree, ast.ExprNode) != 0:
        printf("FAIL: Module misclassified as Stmt/ExprNode\n")
    else:
        printf("PASS: Module is neither Stmt nor ExprNode\n")

    # 断言 2: 顶层 body 每条语句均为 Stmt
    top_ok: t.CInt = 1
    top_n: t.CSizeT = tree.children.__len__()
    i2: t.CSizeT = 0
    while i2 < top_n:
        top_node: ast.AST | t.CPtr = tree.children.get(i2)
        if isinstance(top_node, ast.Stmt) == 0:
            top_ok = 0
            printf("FAIL: top-level node not Stmt: %s\n", top_node.type_name())
        i2 += 1
    if top_ok != 0:
        printf("PASS: all %lld top-level nodes are Stmt\n", top_n)

    # 断言 3: Assign 是 Stmt 且不是 ExprNode
    first_assign: ast.Assign | t.CPtr = (ast.Assign | t.CPtr)(tree.children.get(0))
    if isinstance(first_assign, ast.Stmt) != 0 and isinstance(first_assign, ast.ExprNode) == 0:
        printf("PASS: Assign is Stmt, not ExprNode\n")
    else:
        printf("FAIL: Assign classification wrong\n")

    # 断言 4: Assign.value（表达式）是 ExprNode 且不是 Stmt
    if isinstance(first_assign.value, ast.ExprNode) != 0 \
            and isinstance(first_assign.value, ast.Stmt) == 0:
        printf("PASS: Assign.value is ExprNode (%s)\n", first_assign.value.type_name())
    else:
        printf("FAIL: Assign.value not ExprNode\n")

    # 断言 5: 表达式语句 Expr（如 print(x)）属于 Stmt（Statement 语义）
    if top_n >= 15:
        if_node: ast.If | t.CPtr = (ast.If | t.CPtr)(tree.children.get(14))
        if if_node.children is not None and if_node.children.__len__() > 0:
            expr_stmt: ast.Expr | t.CPtr = (ast.Expr | t.CPtr)(if_node.children.get(0))
            if isinstance(expr_stmt, ast.Stmt) != 0 and isinstance(expr_stmt, ast.ExprNode) == 0:
                printf("PASS: Expr(stmt) is Stmt, not ExprNode (%s)\n", expr_stmt.type_name())
            else:
                printf("FAIL: Expr(stmt) classification wrong\n")
    printf("=== END RUNTIME TEST ===\n")

    stdlib.free(pool_mem)
    stdlib.free(dump_mem)
    return 0