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 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")

    # 解析
    tree: ast.AST | t.CPtr = ast.parse(src, pool)
    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)

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