import stdio
import t, c
import memhub
import stdlib
import string


# ============================================================
# memhub_test - 内存管理库测试
#
# 测试 MemPool/MemBuddy 的 alloc/free/reset/realloc，
# 以及 MemManager 多态虚分派（通过 __vtable__ 调用子类方法）。
# ============================================================


def memhub_test() -> int:
    stdio.printf("memhub: === Test Start ===\n")

    # === Test 1: MemPool (bump allocator) ===
    # MemPool 是 arena bump 分配器，不支持单个 free，reset 一次性回收
    stdio.printf("memhub: === Test 1: MemPool ===\n")

    arena1: bytes = stdlib.malloc(4096)
    if arena1 is None:
        stdio.printf("[FAIL] malloc arena1\n")
        return 1

    pool: memhub.MemPool | t.CPtr = memhub.MemPool(arena1, 4096)

    p1: bytes = pool.alloc(64)
    if p1 is None:
        stdio.printf("[FAIL] MemPool.alloc(64)\n")
        return 1

    p2: bytes = pool.alloc(128)
    if p2 is None:
        stdio.printf("[FAIL] MemPool.alloc(128)\n")
        return 1

    stdio.printf("memhub: MemPool alloc ok\n")

    # 写入验证
    string.memset(p1, 66, 64)
    if p1[0] != 66:
        stdio.printf("[FAIL] MemPool write/read p1[0]=%d\n", p1[0])
        return 1

    # reset 后可以重新分配
    pool.reset()
    p3: bytes = pool.alloc(64)
    if p3 is None:
        stdio.printf("[FAIL] MemPool.alloc after reset\n")
        return 1
    stdio.printf("memhub: MemPool reset ok\n")

    # === Test 2: MemBuddy (buddy system) ===
    # MemBuddy 是伙伴系统，支持 alloc/free/realloc 合并
    stdio.printf("memhub: === Test 2: MemBuddy ===\n")

    arena2: bytes = stdlib.malloc(65536)
    if arena2 is None:
        stdio.printf("[FAIL] malloc arena2\n")
        return 1

    buddy: memhub.MemBuddy | t.CPtr = memhub.MemBuddy(arena2, 65536)

    b1: bytes = buddy.alloc(64)
    if b1 is None:
        stdio.printf("[FAIL] MemBuddy.alloc(64)\n")
        return 1

    b2: bytes = buddy.alloc(256)
    if b2 is None:
        stdio.printf("[FAIL] MemBuddy.alloc(256)\n")
        return 1
    stdio.printf("memhub: MemBuddy alloc ok\n")

    # 写入验证
    string.memset(b1, 85, 64)
    if b1[0] != 85:
        stdio.printf("[FAIL] MemBuddy write/read b1[0]=%d\n", b1[0])
        return 1

    # free b1 后重新 alloc，应能回收内存
    buddy.free(b1)
    b3: bytes = buddy.alloc(64)
    if b3 is None:
        stdio.printf("[FAIL] MemBuddy.alloc after free\n")
        return 1
    stdio.printf("memhub: MemBuddy free/alloc ok\n")

    # === Test 3: MemBuddy realloc ===
    # realloc 扩展内存，验证数据保持
    stdio.printf("memhub: === Test 3: MemBuddy realloc ===\n")

    b4: bytes = buddy.alloc(32)
    if b4 is None:
        stdio.printf("[FAIL] MemBuddy.alloc(32) for realloc\n")
        return 1
    string.memset(b4, 171, 32)

    # realloc 到更大（32 -> 128）
    b4_new: bytes = buddy.realloc(b4, 32, 128)
    if b4_new is None:
        stdio.printf("[FAIL] MemBuddy.realloc\n")
        return 1

    # 验证数据保持（前32字节应保持 171=0xAB）
    # 注意：b4_new[0] 是 i8（有符号），0xAB=-85，需 & 255 转无符号再比较
    if (b4_new[0] & 255) != 171:
        stdio.printf("[FAIL] MemBuddy.realloc data lost b4_new[0]=%d\n", b4_new[0])
        return 1
    stdio.printf("memhub: MemBuddy realloc ok\n")

    # === Test 4: Polymorphism (MemManager virtual dispatch) ===
    # 用 MemManager 指针调用 MemBuddy 的 alloc，验证虚分派
    stdio.printf("memhub: === Test 4: Polymorphism ===\n")

    mgr: memhub.MemManager | t.CPtr = (memhub.MemManager | t.CPtr)(buddy)
    if mgr is None:
        stdio.printf("[FAIL] MemManager cast\n")
        return 1

    # 通过 MemManager 指针调用 alloc（虚分派到 MemBuddy.alloc）
    b5: bytes = mgr.alloc(64)
    if b5 is None:
        stdio.printf("[FAIL] MemManager.alloc (virtual dispatch)\n")
        return 1

    # 验证写入正常
    string.memset(b5, 99, 64)
    if b5[0] != 99:
        stdio.printf("[FAIL] MemManager.alloc write/read\n")
        return 1
    stdio.printf("memhub: Polymorphism ok\n")

    stdio.printf("memhub: === All Tests Passed ===\n")
    return 0