import t, c
from stdint import *
import string
import atom

# 伙伴系统内存分配器
# 支持可变大小分配 (malloc/free/calloc/realloc)
# 通过二分块 (buddy splitting) 和合并 (coalescing) 管理内存
# 使用 atom 库的原子操作实现自旋锁 (线程安全)

MBUDDY_MIN_BLOCK: t.CDefine = 32   # 最小块大小 (字节)，需 >= 16 以容纳链表指针
MBUDDY_MAX_ORDERS: t.CDefine = 32  # 最大阶数数量
MBUDDY_HEADER_SIZE: t.CDefine = 8  # 块头大小 (存储阶数)


def _largest_pow2_le(val: t.CSizeT) -> t.CSizeT:
    """返回 <= val 的最大 2 的幂"""
    if val == 0: return 0
    p: t.CSizeT = 1
    while p * 2 <= val:
        p = p * 2
    return p


def _block_size_at_order(order: t.CInt) -> t.CSizeT:
    """返回指定阶数的块大小"""
    bs: t.CSizeT = MBUDDY_MIN_BLOCK
    i: t.CInt
    for i in range(order):
        bs = bs << 1
    return bs


class MBuddy:
    # 编译期元数据：声明此类在 with 上下文中提供 __mbuddy__ 字段
    __provides__: list[str] = ['__mbuddy__']
    mem: t.CVoid | t.CPtr            # 可用内存基址
    mem_size: t.CSizeT              # 可用内存大小 (2 的幂)
    max_order: t.CInt               # 最大阶数
    free_lists: t.CUInt64T | t.CPtr # 空闲链头数组 (存储在 arena 开头)
    lock_val: t.CVolatile | t.CInt   # 自旋锁标志 (0=未锁, 1=已锁)

    def __init__(self, arena: t.CVoid | t.CPtr, arena_size: t.CSizeT):
        # 初始化自旋锁
        self.lock_val = 0
        # 空闲链数组放在 arena 开头
        fl_bytes: t.CSizeT = (MBUDDY_MAX_ORDERS + 1) * 8
        self.free_lists = arena

        # 初始化所有空闲链头为 0 (NULL)
        i: t.CInt
        for i in range(MBUDDY_MAX_ORDERS + 1):
            self.free_lists[i] = 0

        # 计算可用区域
        if arena_size <= fl_bytes:
            self.mem = None
            self.mem_size = 0
            self.max_order = 0
            return

        remaining: t.CSizeT = arena_size - fl_bytes
        usable: t.CSizeT = _largest_pow2_le(remaining)

        if usable < MBUDDY_MIN_BLOCK:
            self.mem = None
            self.mem_size = 0
            self.max_order = 0
            return

        self.mem = t.CVoid(t.CUInt64T(arena) + fl_bytes, t.CPtr)
        self.mem_size = usable

        # 计算 max_order: log2(usable / MIN_BLOCK)
        self.max_order = 0
        bs: t.CSizeT = MBUDDY_MIN_BLOCK
        while bs < usable:
            bs = bs << 1
            self.max_order += 1

        # 将整个可用区域作为一个 max_order 阶的空闲块加入空闲链
        self._fl_push(self.max_order, self.mem)

    def __enter__(self) -> 'MBuddy' | t.CPtr:
        return self

    def __exit__(self):
        self.reset()

    def reset(self):
        """重置分配器: 清空所有空闲链，将整个可用区域作为单个大块归还"""
        if self.mem == None:
            return
        # 清空所有空闲链头
        i: t.CInt
        for i in range(MBUDDY_MAX_ORDERS + 1):
            self.free_lists[i] = 0
        # 将整个可用区域作为单个 max_order 阶块加入
        self._fl_push(self.max_order, self.mem)

    # === 空闲链操作 ===

    def _fl_push(self, order: t.CInt, block: t.CVoid | t.CPtr):
        """将块压入指定阶数的空闲链头部"""
        old_head: t.CUInt64T = self.free_lists[order]
        c.DerefAs(block, t.CVoid(old_head, t.CPtr))
        self.free_lists[order] = t.CUInt64T(block)

    def _fl_pop(self, order: t.CInt) -> t.CVoid | t.CPtr:
        """从指定阶数的空闲链头部弹出块"""
        head_val: t.CUInt64T = self.free_lists[order]
        if head_val == 0:
            return None
        block: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
        next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        self.free_lists[order] = t.CUInt64T(next_ptr)
        return block

    def _fl_find_and_remove(self, order: t.CInt, target: t.CVoid | t.CPtr) -> t.CInt:
        """在指定阶数的空闲链中查找并移除 target 块。返回 1=找到, 0=未找到"""
        head_val: t.CUInt64T = self.free_lists[order]
        if head_val == 0:
            return 0

        head: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
        if t.CUInt64T(head) == t.CUInt64T(target):
            next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(head, t.CPtr)), t.CPtr)
            self.free_lists[order] = t.CUInt64T(next_ptr)
            return 1

        prev: t.CVoid | t.CPtr = head
        cur: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(head, t.CPtr)), t.CPtr)
        while cur != None:
            if t.CUInt64T(cur) == t.CUInt64T(target):
                next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
                c.DerefAs(prev, next_ptr)
                return 1
            prev = cur
            cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)

        return 0

    # === 伙伴系统核心操作 ===

    def _buddy_of(self, block: t.CVoid | t.CPtr, order: t.CInt) -> t.CVoid | t.CPtr:
        """计算指定块的伙伴块地址"""
        offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.mem)
        bs: t.CSizeT = _block_size_at_order(order)
        buddy_offset: t.CSizeT = offset ^ bs
        return t.CVoid(t.CUInt64T(self.mem) + buddy_offset, t.CPtr)

    def _order_for_size(self, size: t.CSizeT) -> t.CInt:
        """返回使 MIN_BLOCK << order >= size 的最小阶数"""
        order: t.CInt = 0
        bs: t.CSizeT = MBUDDY_MIN_BLOCK
        while bs < size:
            bs = bs << 1
            order += 1
        return order

    def _split_to_order(self, to_order: t.CInt) -> t.CVoid | t.CPtr:
        """找到可用块并分裂到目标阶数。返回块指针或 None"""
        # 寻找第一个有空闲块的阶数 >= to_order
        found_order: t.CInt = to_order
        while found_order <= self.max_order:
            if self.free_lists[found_order] != 0:
                break
            found_order += 1

        if found_order > self.max_order:
            return None

        # 弹出块
        block: t.CVoid | t.CPtr = self._fl_pop(found_order)

        # 逐级分裂到 to_order
        while found_order > to_order:
            found_order -= 1
            bs: t.CSizeT = _block_size_at_order(found_order)
            buddy: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(block) + bs, t.CPtr)
            self._fl_push(found_order, buddy)

        return block

    def _coalesce(self, block: t.CVoid | t.CPtr, order: t.CInt):
        """尝试与伙伴合并，递归直到无法合并"""
        while order < self.max_order:
            buddy: t.CVoid | t.CPtr = self._buddy_of(block, order)
            found: t.CInt = self._fl_find_and_remove(order, buddy)
            if found == 0:
                break
            # 使用地址较低的一半作为合并后的块
            if t.CUInt64T(buddy) < t.CUInt64T(block):
                block = buddy
            order += 1

        self._fl_push(order, block)

    def _is_valid_ptr(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        """验证指针是否为有效的已分配块指针。返回 1=有效, 0=无效"""
        if ptr == None:
            return 0
        if self.mem == None:
            return 0
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_HEADER_SIZE, t.CPtr)
        # 检查范围
        if t.CUInt64T(block) < t.CUInt64T(self.mem):
            return 0
        if t.CUInt64T(block) >= t.CUInt64T(self.mem) + self.mem_size:
            return 0
        # 检查对齐 (所有块都 MIN_BLOCK 对齐)
        offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.mem)
        if offset % MBUDDY_MIN_BLOCK != 0:
            return 0
        # 读取头部编码值: (order << 1) | allocated_flag
        stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        stored_val: t.CUInt64T = t.CUInt64T(stored)
        # 检查 allocated 标志位 (bit 0 = 1 表示已分配)
        if (stored_val & 1) == 0:
            return 0  # 未分配或已释放 (双重 free 检测)
        # 提取阶数
        order: t.CInt = t.CInt(stored_val >> 1)
        if order < 0:
            return 0
        if order > self.max_order:
            return 0
        return 1

    def _fl_count(self, order: t.CInt) -> t.CSizeT:
        """统计指定阶数空闲链中的块数"""
        count: t.CSizeT = 0
        cur_val: t.CUInt64T = self.free_lists[order]
        while cur_val != 0:
            count += 1
            cur: t.CVoid | t.CPtr = t.CVoid(cur_val, t.CPtr)
            cur_val = t.CUInt64T(t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr))
        return count

    # === 公共 API ===

    def _lock(self):
        """自旋锁: 原子测试并设置 lock_val，旧值为 1 时自旋等待"""
        while atom.__atomic_test_and_set(c.Addr(self.lock_val), atom.ATOMIC_ACQUIRE):
            pass

    def _unlock(self):
        """解锁: 原子清除 lock_val"""
        atom.__atomic_clear(c.Addr(self.lock_val), atom.ATOMIC_RELEASE)

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        """分配至少 size 字节的内存。返回用户数据指针或 None (线程安全)"""
        self._lock()
        result: t.CVoid | t.CPtr = None
        if self.mem != None:
            if size != 0:
                needed: t.CSizeT = size + MBUDDY_HEADER_SIZE
                order: t.CInt = self._order_for_size(needed)
                if order <= self.max_order:
                    block: t.CVoid | t.CPtr = self._split_to_order(order)
                    if block != None:
                        c.DerefAs(block, t.CVoid(t.CUInt64T((order << 1) | 1), t.CPtr))
                        result = t.CVoid(t.CUInt64T(block) + MBUDDY_HEADER_SIZE, t.CPtr)
        self._unlock()
        return result

    def free(self, ptr: t.CVoid | t.CPtr):
        """释放之前分配的内存。非法指针将被忽略 (线程安全)"""
        self._lock()
        if ptr != None:
            if self._is_valid_ptr(ptr) != 0:
                # 从头部读取阶数 (编码: (order << 1) | 1)
                block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_HEADER_SIZE, t.CPtr)
                stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
                stored_val: t.CUInt64T = t.CUInt64T(stored)
                order: t.CInt = t.CInt(stored_val >> 1)
                # 清除 allocated 标志，防止双重 free
                c.DerefAs(block, t.CVoid(0, t.CPtr))
                # 合并并加入空闲链
                self._coalesce(block, order)
        self._unlock()

    def calloc(self, count: t.CSizeT, size: t.CSizeT) -> t.CVoid | t.CPtr:
        """分配 count*size 字节并清零"""
        total: t.CSizeT = count * size
        ptr: t.CVoid | t.CPtr = self.alloc(total)
        if ptr != None:
            string.memset(ptr, 0, total)
        return ptr

    def realloc(self, ptr: t.CVoid | t.CPtr, new_size: t.CSizeT) -> t.CVoid | t.CPtr:
        """重新分配内存大小"""
        if ptr == None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        if self._is_valid_ptr(ptr) == 0:
            return None  # 非法指针

        # 读取旧阶数 (编码: (order << 1) | 1)
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_HEADER_SIZE, t.CPtr)
        stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        stored_val: t.CUInt64T = t.CUInt64T(stored)
        old_order: t.CInt = t.CInt(stored_val >> 1)

        # 计算新阶数
        needed: t.CSizeT = new_size + MBUDDY_HEADER_SIZE
        new_order: t.CInt = self._order_for_size(needed)

        # 如果新阶数 <= 旧阶数，无需重新分配
        if new_order <= old_order:
            return ptr

        # 分配新块
        new_ptr: t.CVoid | t.CPtr = self.alloc(new_size)
        if new_ptr == None:
            return ptr  # 保留旧块

        # 复制旧数据
        old_block_size: t.CSizeT = _block_size_at_order(old_order)
        old_data_size: t.CSizeT = old_block_size - MBUDDY_HEADER_SIZE
        string.memcpy(new_ptr, ptr, old_data_size)

        # 释放旧块
        self.free(ptr)
        return new_ptr

    def stats(self) -> t.CSizeT:
        """返回总可用内存大小"""
        return self.mem_size

    def free_count(self) -> t.CSizeT:
        """统计所有空闲链中的总块数"""
        total: t.CSizeT = 0
        i: t.CInt
        for i in range(self.max_order + 1):
            total += self._fl_count(i)
        return total

    def self_check(self) -> t.CInt:
        """自检: 遍历所有空闲链，验证块的范围、对齐和阶数一致性。返回错误数 (0=正常)"""
        if self.mem == None:
            return 0
        errors: t.CInt = 0
        order: t.CInt
        for order in range(self.max_order + 1):
            bs: t.CSizeT = _block_size_at_order(order)
            max_in_order: t.CSizeT = self.mem_size / bs + 1
            seen: t.CSizeT = 0
            cur_val: t.CUInt64T = self.free_lists[order]
            while cur_val != 0:
                seen += 1
                if seen > max_in_order:
                    errors += 1
                    break
                cur: t.CVoid | t.CPtr = t.CVoid(cur_val, t.CPtr)
                # 检查范围
                if t.CUInt64T(cur) < t.CUInt64T(self.mem):
                    errors += 1
                    break
                if t.CUInt64T(cur) >= t.CUInt64T(self.mem) + self.mem_size:
                    errors += 1
                    break
                # 检查对齐
                offset: t.CSizeT = t.CUInt64T(cur) - t.CUInt64T(self.mem)
                if offset % MBUDDY_MIN_BLOCK != 0:
                    errors += 1
                    break
                # 检查阶数一致性: 块偏移应是其块大小的倍数
                if offset % bs != 0:
                    errors += 1
                    break
                cur_val = t.CUInt64T(t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr))
        return errors