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 size: 1MB
define ARENA_SIZE = 1048576
# mpool memory size: 512KB
define POOL_SIZE = 524288
# dump buffer size: 256KB
define DUMP_SIZE = 262144


def classify_list(lst: list[ast.AST*]*, depth: int):
    """Recursively walk a list[AST | t.CPtr] container, classify each element."""
    if lst is None:
        return
    n: sizet = lst.__len__()
    for i in range(n):
        classify_print(lst.get(i), depth + 1)


def classify_print(node: ast.AST*, depth: int):
    """Recursively walk the parse tree and print each node's category at runtime
    via TPV isinstance vtable subclass-set comparison (Stmt vs ExprNode vs other).
    """
    if node is None:
        return
    for sp in range(depth):
        printf("  ")
    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())
    # Walk statement expression fields by kind (bodies live in children); cast to
    # the concrete class first, since ast.AST base-class field access fast-fails.
    k: int = node.kind()
    if k == ast.ASTKind.Assign:
        asg: ast.Assign* = (ast.Assign*)(node)
        classify_print(asg.value, depth + 1)
    elif k == ast.ASTKind.If:
        ifn: ast.If* = (ast.If*)(node)
        classify_print(ifn.test, depth + 1)
        classify_list(ifn.orelse, depth + 1)
    elif k == ast.ASTKind.While:
        whl: ast.While* = (ast.While*)(node)
        classify_print(whl.test, depth + 1)
        classify_list(whl.orelse, depth + 1)
    elif k == ast.ASTKind.For:
        for_nd: ast.For* = (ast.For*)(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* = (ast.Return*)(node)
        classify_print(ret_nd.value, depth + 1)
    elif k == ast.ASTKind.Expr:
        expr_nd: ast.Expr* = (ast.Expr*)(node)
        classify_print(expr_nd.value, depth + 1)
    elif k == ast.ASTKind.BinOp:
        bin_nd: ast.BinOp* = (ast.BinOp*)(node)
        classify_print(bin_nd.left, depth + 1)
        classify_print(bin_nd.right, depth + 1)
    elif k == ast.ASTKind.UnaryOp:
        una_nd: ast.UnaryOp* = (ast.UnaryOp*)(node)
        classify_print(una_nd.operand, depth + 1)
    elif k == ast.ASTKind.Compare:
        cmp_nd: ast.Compare* = (ast.Compare*)(node)
        classify_print(cmp_nd.left, depth + 1)
        classify_list(cmp_nd.comparators, depth + 1)
    elif k == ast.ASTKind.BoolOp:
        bool_nd: ast.BoolOp* = (ast.BoolOp*)(node)
        classify_list(bool_nd.values, depth + 1)
    elif k == ast.ASTKind.Call:
        call_nd: ast.Call* = (ast.Call*)(node)
        classify_print(call_nd.func, depth + 1)
        classify_list(call_nd.args, depth + 1)
    else:
        pass
    # children (statement bodies)
    if node.children is not None:
        n: sizet = node.children.__len__()
        for i in range(n):
            classify_print(node.children.get(i), depth + 1)


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

    printf("AST Comprehensive Test Start\n")

    # Initialize the global MBuddy allocator
    arena: bytes = stdlib.malloc(ARENA_SIZE)
    if arena == None:
        printf("FAIL: arena alloc failed\n")
        return 1
    mbuddy_obj: memhub.MemBuddy* = 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

    # Create an mpool for the ast parser
    pool_mem: bytes = stdlib.malloc(POOL_SIZE)
    if pool_mem == None:
        printf("FAIL: malloc for pool failed\n")
        return 1
    pool: memhub.MemPool* = memhub.MemPool(pool_mem, POOL_SIZE)
    if pool == None:
        printf("FAIL: memhub.MemPool init failed\n")
        return 1

    # dump buffer (separately allocated)
    dump_mem: bytes = stdlib.malloc(DUMP_SIZE)
    if dump_mem == None:
        printf("FAIL: malloc for dump failed\n")
        return 1
    dump_buf: char* = char(dump_mem, t.CPtr)

    # Comprehensive test source covering main CPython syntax
    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")

    # Parse
    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* = 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 runtime distinction test (isinstance vtable compare) =====
    printf("=== STMT/EXPR RUNTIME DISTINGUISH TEST ===\n")
    classify_print(tree, 0)

    # Assert 1: Module is neither Stmt nor ExprNode (top-level container)
    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")

    # Assert 2: every top-level body statement is a Stmt
    top_ok: int = 1
    top_n: sizet = tree.children.__len__()
    for i2 in range(top_n):
        top_node: ast.AST* = 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())
    if top_ok != 0:
        printf("PASS: all %lld top-level nodes are Stmt\n", top_n)

    # Assert 3: Assign is Stmt, not ExprNode
    first_assign: ast.Assign* = (ast.Assign*)(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")

    # Assert 4: Assign.value (an expression) is ExprNode, not 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")

    # Assert 5: expression statement Expr (e.g. print(x)) is a Stmt
    if top_n >= 15:
        if_node: ast.If* = (ast.If*)(tree.children.get(14))
        if if_node.children is not None and if_node.children.__len__() > 0:
            expr_stmt: ast.Expr* = (ast.Expr*)(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