import viperlib
import viperio
import stdlib
import memhub
import string
import testcheck
import win32.win32process
import win32.win32sync
import win32.win32base


# Thread body: alloc/free on a shared MBuddy to exercise the spinlock
def _thread_alloc_free(param: rawptr) -> uint32:
    bd: memhub.MemBuddy* = (memhub.MemBuddy*)(param)
    i: int
    for i in range(100):
        p: bytes = bd.alloc(32)
        if p != None:
            p_c: uint8* = uint8*(p)
            p_c[0] = uint8(i & 0xFF)
            bd.free(p)
    return 0


def mbuddy_main() -> int:
    buf: bytes = stdlib.malloc(256)

    testcheck.begin("BuddyTest: 伙伴系统内存分配器测试")

    # === Basic alloc/free ===
    testcheck.section("基础 alloc/free")
    arena: bytes = stdlib.malloc(65536)
    bd: memhub.MemBuddy* = memhub.MemBuddy(arena, 65536)

    viperlib.snprintf(buf, 256, "stats=%lu, free_count=%lu", bd.stats(), bd.free_count())
    testcheck.info(buf)

    p1: bytes = bd.alloc(16)
    p2: bytes = bd.alloc(64)
    p3: bytes = bd.alloc(256)
    testcheck.check(p1 != None and p2 != None and p3 != None,
                    "3 allocs OK", "alloc FAILED")

    # Write data to verify the returned pointer is usable
    p_c: uint8* = uint8*(p1)
    p_c[0] = uint8(0xAB)
    p_c[1] = uint8(0xCD)
    viperlib.snprintf(buf, 256, "p1[0]=0x%x p1[1]=0x%x", p_c[0], p_c[1])
    testcheck.info(buf)
    testcheck.check(p_c[0] == uint8(0xAB) and p_c[1] == uint8(0xCD),
                    "write/read OK", "write/read FAILED")

    bd.free(p1)
    bd.free(p2)
    bd.free(p3)
    testcheck.ok("3 frees OK")

    # After free, coalescing must restore the initial 10 free blocks
    # (65536 - 264 = 65272 decomposes into 10 power-of-two blocks)
    testcheck.check(bd.free_count() == 10,
                    "coalesce: free_count=10 OK", "coalesce: free_count != 10")

    # Coalesced space must allow a large alloc again
    big: bytes = bd.alloc(16000)
    testcheck.check(big != None, "big alloc OK", "big alloc FAILED")
    bd.free(big)

    stdlib.free(arena)

    # === calloc (verify zeroing) ===
    testcheck.section("calloc 清零")
    arena2: bytes = stdlib.malloc(4096)
    viperlib.snprintf(buf, 256, "arena2 = %p", arena2)
    testcheck.info(buf)
    bd2: memhub.MemBuddy* = memhub.MemBuddy(arena2, 4096)

    dirty: bytes = bd2.alloc(64)
    viperlib.snprintf(buf, 256, "dirty = %p", dirty)
    testcheck.info(buf)
    d_c: uint8* = uint8*(dirty)
    i: int
    for i in range(64):
        d_c[i] = uint8(0xFF)
    bd2.free(dirty)

    cl: bytes = bd2.calloc(8, 8)
    viperlib.snprintf(buf, 256, "cl = %p (arena2=%p, dirty was=%p)", cl, arena2, dirty)
    testcheck.info(buf)
    if cl != None:
        cl_c: uint8* = uint8*(cl)
        all_zero: int = 1
        j: int
        for j in range(64):
            if cl_c[j] != 0:
                all_zero = 0
                break
        testcheck.check(all_zero == 1, "calloc zeroed OK", "calloc NOT zeroed")
        bd2.free(cl)
    else:
        testcheck.fail("calloc FAILED")

    stdlib.free(arena2)

    # === realloc grow ===
    testcheck.section("realloc 扩展")
    arena3: bytes = stdlib.malloc(4096)
    bd3: memhub.MemBuddy* = memhub.MemBuddy(arena3, 4096)

    rp: bytes = bd3.alloc(32)
    rp_c: uint8* = uint8*(rp)
    k: int
    for k in range(32):
        rp_c[k] = uint8(k + 1)

    rp2: bytes = bd3.realloc(rp, 128)
    if rp2 != None:
        rp2_c: uint8* = uint8*(rp2)
        ok: int = 1
        m: int
        for m in range(32):
            if rp2_c[m] != uint8(m + 1):
                ok = 0
                break
        testcheck.check(ok == 1, "realloc grow data preserved OK", "realloc grow data LOST")
        bd3.free(rp2)
    else:
        testcheck.fail("realloc grow FAILED")

    stdlib.free(arena3)

    # === in-place shrink (verify no leak) ===
    testcheck.section("原地收缩 (realloc shrink)")
    arena_shrink: bytes = stdlib.malloc(4096)
    bds: memhub.MemBuddy* = memhub.MemBuddy(arena_shrink, 4096)

    sp: bytes = bds.alloc(256)
    testcheck.check(sp != None, "alloc 256 OK", "alloc 256 FAILED")

    fc_before: sizet = bds.free_count()
    # Shrink 256 -> 32: order 4 (264) to order 1 (40); since new order <= old,
    # realloc must return the same pointer in place
    sp2: bytes = bds.realloc(sp, 32)
    fc_after: sizet = bds.free_count()

    testcheck.check(sp2 == sp, "shrink same ptr OK", "shrink ptr changed")
    testcheck.check(fc_after == fc_before, "shrink no leak OK", "shrink LEAKED blocks")

    # Data must remain readable/writable after shrink
    sp_c: uint8* = uint8*(sp2)
    sp_c[0] = uint8(0x42)
    testcheck.check(sp_c[0] == uint8(0x42), "shrink data OK", "shrink data FAILED")

    bds.free(sp2)
    testcheck.check(bds.free_count() == 6, "shrink+free coalesce OK", "shrink+free coalesce FAILED")

    stdlib.free(arena_shrink)

    # === Edge cases ===
    testcheck.section("边界情况")
    arena4: bytes = stdlib.malloc(4096)
    bd4: memhub.MemBuddy* = memhub.MemBuddy(arena4, 4096)

    # alloc(0) must return None
    z: bytes = bd4.alloc(0)
    testcheck.check(z == None, "alloc(0)=None OK", "alloc(0) should be None")

    # free(None) must be safe
    bd4.free(None)
    testcheck.ok("free(None) safe")

    # realloc(None, size) behaves like alloc
    rn: bytes = bd4.realloc(None, 32)
    testcheck.check(rn != None, "realloc(None)=alloc OK", "realloc(None) FAILED")
    bd4.free(rn)

    # realloc(ptr, 0) behaves like free
    rz: bytes = bd4.alloc(32)
    rz2: bytes = bd4.realloc(rz, 0)
    testcheck.check(rz2 == None, "realloc(ptr,0)=free OK", "realloc(ptr,0) should be None")

    stdlib.free(arena4)

    # === Illegal pointer defense ===
    testcheck.section("非法指针防护")
    arena5: bytes = stdlib.malloc(4096)
    bd5: memhub.MemBuddy* = memhub.MemBuddy(arena5, 4096)

    valid: bytes = bd5.alloc(32)
    testcheck.check(valid != None, "alloc for illegal test OK", "alloc FAILED")

    # 1. free a random address (address 1)
    bd5.free(rawptr(uint64(1)))
    testcheck.ok("free(addr=1) ignored")

    # 2. free NULL
    bd5.free(None)
    testcheck.ok("free(NULL) ignored")

    # 3. free an out-of-range pointer (before arena)
    bd5.free(rawptr(uint64(arena5)))
    testcheck.ok("free(before arena) ignored")

    # 4. free an out-of-range pointer (after arena)
    bd5.free(rawptr(uint64(arena5) + 99999))
    testcheck.ok("free(after arena) ignored")

    # 5. free a misaligned pointer (inside mem but not aligned)
    bd5.free(rawptr(uint64(bd5.base) + 1))
    testcheck.ok("free(misaligned) ignored")

    # 6. a valid free must still work (block reusable after)
    bd5.free(valid)
    retry5: bytes = bd5.alloc(32)
    testcheck.check(retry5 != None, "valid free still works OK", "valid free BROKEN")
    if retry5 != None:
        bd5.free(retry5)

    # 7. realloc on an illegal pointer must return None
    r_illegal: bytes = bd5.realloc(rawptr(1), 64)
    testcheck.check(r_illegal == None, "realloc(illegal)=None OK", "realloc(illegal) should be None")

    stdlib.free(arena5)

    # === OOM exhaustion ===
    testcheck.section("OOM 内存耗尽")
    arena6: bytes = stdlib.malloc(4096)
    bd6: memhub.MemBuddy* = memhub.MemBuddy(arena6, 4096)

    # Allocate until failure
    alloc_count: int = 0
    ptrs_arr: bytes = stdlib.malloc(8 * 256)  # Pointer storage array
    ptrs: uint64* = uint64*(ptrs_arr)

    while alloc_count < 256:
        p: bytes = bd6.alloc(32)
        if p == None:
            break
        ptrs[alloc_count] = uint64(p)
        alloc_count += 1

    viperlib.snprintf(buf, 256, "allocated %d blocks before OOM", alloc_count)
    testcheck.info(buf)
    testcheck.check(alloc_count > 0, "OOM: got some allocs OK", "OOM: no allocs at all")

    # Another alloc must now fail
    oom_p: bytes = bd6.alloc(32)
    testcheck.check(oom_p == None, "OOM: alloc fails OK", "OOM: alloc should fail")

    # After freeing one block, allocation must succeed again
    if alloc_count > 0:
        first_ptr: bytes = rawptr(ptrs[0])
        bd6.free(first_ptr)
        retry: bytes = bd6.alloc(32)
        testcheck.check(retry != None, "OOM: free+alloc OK", "OOM: free+alloc FAILED")
        if retry != None:
            bd6.free(retry)

    # Free all blocks
    for idx in range(1, alloc_count):
        cur_p: bytes = rawptr(ptrs[idx])
        bd6.free(cur_p)

    # Full free must coalesce back to the initial 6 blocks
    testcheck.check(bd6.free_count() == 6,
                    "OOM: full coalesce OK", "OOM: free_count != 6 after full free")

    stdlib.free(ptrs_arr)
    stdlib.free(arena6)

    # === Free-list state ===
    testcheck.section("空闲链表状态")
    arena7: bytes = stdlib.malloc(4096)
    bd7: memhub.MemBuddy* = memhub.MemBuddy(arena7, 4096)

    # Initial: non-power-of-two arena (4096 - 264 = 3832) splits into 6 blocks
    testcheck.check(bd7.free_count() == 6, "init: 6 free blocks OK", "init: free_count != 6")

    # After alloc: free count grows because a block is split
    a1: bytes = bd7.alloc(32)
    fc1: sizet = bd7.free_count()
    viperlib.snprintf(buf, 256, "after alloc(32): free_count=%lu", fc1)
    testcheck.info(buf)
    testcheck.check(fc1 > 1, "after alloc: split OK", "after alloc: no split")

    # Second alloc
    a2: bytes = bd7.alloc(32)
    fc2: sizet = bd7.free_count()
    testcheck.check(fc2 >= 1, "after 2nd alloc: free >= 1 OK", "after 2nd alloc: free < 1")

    # Free a1: free count must increase
    bd7.free(a1)
    fc3: sizet = bd7.free_count()
    testcheck.check(fc3 > fc2, "after free a1: free increased OK", "after free a1: free not increased")

    # Free a2: coalescing must restore the initial 6 blocks
    bd7.free(a2)
    testcheck.check(bd7.free_count() == 6,
                    "after free all: coalesced to 6 OK", "after free all: free_count != 6")

    stdlib.free(arena7)

    # === with context manager ===
    testcheck.section("with 上下文管理器")
    arena8: bytes = stdlib.malloc(4096)
    with memhub.MemBuddy(arena8, 4096) as wbd:
        wp1: bytes = wbd.alloc(32)
        wp2: bytes = wbd.alloc(64)
        testcheck.check(wp1 != None and wp2 != None,
                        "with: allocs OK", "with: allocs FAILED")
        # Leaving the with block triggers __exit__ auto-reset

    # Re-enter after exit to verify the reset took effect
    with memhub.MemBuddy(arena8, 4096) as wbd2:
        wp3: bytes = wbd2.alloc(1024)
        testcheck.check(wp3 != None, "with: reset+realloc OK", "with: reset+realloc FAILED")

    stdlib.free(arena8)

    # === Double-free detection ===
    testcheck.section("双重 free 检测")
    arena_df: bytes = stdlib.malloc(4096)
    bd_df: memhub.MemBuddy* = memhub.MemBuddy(arena_df, 4096)

    df: bytes = bd_df.alloc(64)
    testcheck.check(df != None, "df: alloc OK", "df: alloc FAILED")

    # First free
    bd_df.free(df)
    fc_df1: sizet = bd_df.free_count()
    testcheck.ok("first free OK")

    # Second free (double free): must be silently rejected
    bd_df.free(df)
    fc_df2: sizet = bd_df.free_count()
    testcheck.check(fc_df2 == fc_df1, "double free rejected OK", "double free CORRUPTED state")

    # Allocator must still be functional
    df2: bytes = bd_df.alloc(64)
    testcheck.check(df2 != None, "df: alloc after double-free OK", "df: alloc after double-free FAILED")
    if df2 != None:
        bd_df.free(df2)

    stdlib.free(arena_df)

    # === Self-check ===
    testcheck.section("自检功能")
    arena_sc: bytes = stdlib.malloc(4096)
    bd_sc: memhub.MemBuddy* = memhub.MemBuddy(arena_sc, 4096)

    # Self-check on initial state
    testcheck.check(bd_sc.self_check() == 0, "self_check init OK", "self_check init FAILED")

    # Self-check after alloc
    sc_p1: bytes = bd_sc.alloc(32)
    sc_p2: bytes = bd_sc.alloc(128)
    testcheck.check(bd_sc.self_check() == 0, "self_check after alloc OK", "self_check after alloc FAILED")

    # Self-check after partial free
    bd_sc.free(sc_p1)
    testcheck.check(bd_sc.self_check() == 0, "self_check after partial free OK", "self_check after partial free FAILED")

    # Self-check after full free (should coalesce to one block)
    bd_sc.free(sc_p2)
    testcheck.check(bd_sc.self_check() == 0, "self_check after full free OK", "self_check after full free FAILED")

    # Self-check after double free (state must be intact)
    sc_p3: bytes = bd_sc.alloc(64)
    bd_sc.free(sc_p3)
    bd_sc.free(sc_p3)  # Double free, must be rejected
    testcheck.check(bd_sc.self_check() == 0, "self_check after double-free OK", "self_check after double-free CORRUPTED")

    stdlib.free(arena_sc)

    # === Concurrent threads (spinlock) ===
    testcheck.section("并发多线程 (spinlock)")
    arena_mt: bytes = stdlib.malloc(65536)
    bd_mt: memhub.MemBuddy* = memhub.MemBuddy(arena_mt, 65536)

    NUM_THREADS: int = 4
    handles_arr: bytes = stdlib.malloc(8 * 4)
    handles: uint64* = uint64*(handles_arr)
    tid: uint32 = 0

    # Spawn 4 threads doing concurrent alloc/free
    t_idx: int
    for t_idx in range(NUM_THREADS):
        h: bytes = win32.win32process.CreateThread(None, 0, _thread_alloc_free, rawptr(bd_mt), 0, &tid)
        handles[t_idx] = uint64(h)

    # Wait for all threads to finish
    for t_idx in range(NUM_THREADS):
        h: bytes = rawptr(handles[t_idx])
        win32.win32sync.WaitForSingleObject(h, win32.win32base.INFINITE)
        win32.win32base.CloseHandle(h)

    # Verify allocator state: self-check passes and all blocks coalesced
    testcheck.check(bd_mt.self_check() == 0, "mt: self_check OK", "mt: self_check FAILED")
    testcheck.check(bd_mt.free_count() == 10, "mt: free_count=10 OK", "mt: free_count != 10 (leaked)")

    viperlib.snprintf(buf, 256, "mt: free_count=%lu", bd_mt.free_count())
    testcheck.info(buf)

    stdlib.free(handles_arr)
    stdlib.free(arena_mt)

    # === Extreme boundaries ===
    testcheck.section("极端边界")

    # 1. Alloc size exactly at a block boundary
    arena_ex: bytes = stdlib.malloc(4096)
    bd_ex: memhub.MemBuddy* = memhub.MemBuddy(arena_ex, 4096)

    # 24 + 16 = 40 > 32 (8 header + 8 sentinel), needs order 1 (64 bytes)
    exact_fit: bytes = bd_ex.alloc(24)
    testcheck.check(exact_fit != None, "exact fit order 0 OK", "exact fit order 0 FAILED")

    # 25 + 16 = 41 > 32, needs order 1 (64 bytes)
    over_fit: bytes = bd_ex.alloc(25)
    testcheck.check(over_fit != None, "over by 1 byte OK", "over by 1 byte FAILED")

    bd_ex.free(exact_fit)
    bd_ex.free(over_fit)
    testcheck.check(bd_ex.free_count() == 6, "boundary: coalesce OK", "boundary: coalesce FAILED")

    # 2. Alloc near the largest block limit (order6 = 2048): usable = 2048 - 16,
    # so alloc(max_usable) exactly fills one order6 block
    max_usable: sizet = 2048 - 16
    viperlib.snprintf(buf, 256, "mem_size=%lu, max_block=2048, max_usable=%lu", bd_ex.mem_size, max_usable)
    testcheck.info(buf)

    # Exactly consumes the order6 block (5 other blocks stay free)
    exact_max: bytes = bd_ex.alloc(max_usable)
    testcheck.check(exact_max != None, "exact max alloc OK", "exact max alloc FAILED")
    testcheck.check(bd_ex.free_count() == 5, "max alloc: 5 free blocks left OK", "max alloc: free_count != 5")

    bd_ex.free(exact_max)
    testcheck.check(bd_ex.free_count() == 6, "max free: coalesce OK", "max free: coalesce FAILED")

    # Alloc the same size again to verify coalesced space is reusable
    exact_max2: bytes = bd_ex.alloc(max_usable)
    testcheck.check(exact_max2 != None, "exact max re-alloc OK", "exact max re-alloc FAILED")
    bd_ex.free(exact_max2)

    # 3. Alloc beyond the arena limit must fail
    over_max: bytes = bd_ex.alloc(bd_ex.mem_size)
    testcheck.check(over_max == None, "over max: fails OK", "over max: should fail")

    over_max2: bytes = bd_ex.alloc(max_usable + 1)
    testcheck.check(over_max2 == None, "max+1: fails OK", "max+1: should fail")

    stdlib.free(arena_ex)

    # 4. Arena too small (only fits the 264-byte free-list overhead)
    arena_tiny: bytes = stdlib.malloc(264)
    bd_tiny: memhub.MemBuddy* = memhub.MemBuddy(arena_tiny, 264)
    testcheck.check(bd_tiny.stats() == 0, "tiny arena: stats=0 OK", "tiny arena: stats != 0")
    tiny_p: bytes = bd_tiny.alloc(32)
    testcheck.check(tiny_p == None, "tiny arena: alloc fails OK", "tiny arena: should fail")
    stdlib.free(arena_tiny)

    # 5. Arena fits exactly one minimal block (264 + 32 = 296)
    arena_one: bytes = stdlib.malloc(296)
    bd_one: memhub.MemBuddy* = memhub.MemBuddy(arena_one, 296)
    testcheck.check(bd_one.stats() == 32, "one block: stats=32 OK", "one block: stats != 32")

    # Minimal block 32 = data 16 + header 8 + sentinel 8; alloc(16) fills it
    one_p: bytes = bd_one.alloc(16)
    testcheck.check(one_p != None, "one block: alloc OK", "one block: alloc FAILED")

    one_p2: bytes = bd_one.alloc(16)
    testcheck.check(one_p2 == None, "one block: OOM OK", "one block: should be OOM")

    bd_one.free(one_p)
    testcheck.check(bd_one.free_count() == 1, "one block: free OK", "one block: free FAILED")
    stdlib.free(arena_one)

    # 6. Power-of-two vs non-power-of-two arena
    arena_pow2: bytes = stdlib.malloc(512)
    bd_pow2: memhub.MemBuddy* = memhub.MemBuddy(arena_pow2, 512)
    # 512 - 264 = 248, decomposes to 128 + 64 + 32 = 224 (24-byte tail dropped)
    testcheck.check(bd_pow2.stats() == 224, "pow2 arena: stats=224 OK", "pow2 arena: stats != 224")
    stdlib.free(arena_pow2)

    arena_non2: bytes = stdlib.malloc(500)
    bd_non2: memhub.MemBuddy* = memhub.MemBuddy(arena_non2, 500)
    # 500 - 264 = 236, decomposes to 128 + 64 + 32 = 224 (12-byte tail dropped)
    testcheck.check(bd_non2.stats() == 224, "non-pow2 arena: stats=224 OK", "non-pow2 arena: stats != 224")
    stdlib.free(arena_non2)

    # 7. Stress: 1000 alloc/free cycles of the same size
    arena_stress: bytes = stdlib.malloc(4096)
    bd_stress: memhub.MemBuddy* = memhub.MemBuddy(arena_stress, 4096)

    s: int
    for s in range(1000):
        sp: bytes = bd_stress.alloc(64)
        if sp != None:
            sp_c: uint8* = uint8*(sp)
            sp_c[0] = uint8(s & 0xFF)
            bd_stress.free(sp)

    testcheck.check(bd_stress.self_check() == 0, "stress: self_check OK", "stress: self_check FAILED")
    testcheck.check(bd_stress.free_count() == 6, "stress: free_count=6 OK", "stress: leaked")
    stdlib.free(arena_stress)

    # 8. Allocate all minimal blocks, then free them all
    arena_fill: bytes = stdlib.malloc(4096)
    bd_fill: memhub.MemBuddy* = memhub.MemBuddy(arena_fill, 4096)

    fill_count: int = 0
    fill_ptrs_arr: bytes = stdlib.malloc(8 * 256)
    fill_ptrs: uint64* = uint64*(fill_ptrs_arr)

    fc: int
    for fc in range(256):
        fp: bytes = bd_fill.alloc(24)
        if fp == None:
            break
        fill_ptrs[fill_count] = uint64(fp)
        fill_count += 1

    viperlib.snprintf(buf, 256, "fill: %d min blocks allocated", fill_count)
    testcheck.info(buf)
    testcheck.check(fill_count > 0, "fill: got blocks OK", "fill: no blocks")
    testcheck.check(bd_fill.self_check() == 0, "fill: self_check OK", "fill: self_check FAILED")

    # Free all
    fi: int
    for fi in range(fill_count):
        fp2: bytes = rawptr(fill_ptrs[fi])
        bd_fill.free(fp2)

    testcheck.check(bd_fill.free_count() == 6, "fill: all freed OK", "fill: leaked after free all")
    testcheck.check(bd_fill.self_check() == 0, "fill: self_check after free OK", "fill: self_check FAILED")

    stdlib.free(fill_ptrs_arr)
    stdlib.free(arena_fill)

    stdlib.free(buf)
    return testcheck.end()