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

# ============================================================
# memhub - 统一内存管理库
# MemManager (多态基类, vtable)
#   ├─ MemPool  (arena bump 分配器, 不支持单个 free)
#   ├─ MemSlab  (定长块池, 支持 alloc/free 位图跟踪)
#   └─ MemBuddy (伙伴系统, 支持 alloc/free/realloc 合并)
# 所有子类通过 vtable 覆盖 alloc/free/reset, 可用 MemManager 指针多态调用
# ============================================================

MEMHUB_ALIGN: t.CDefine = 8

# MemSlab 位图常量
MEMSLAB_MIN_BLOCK: t.CDefine = 16
MEMSLAB_BITMAP_BYTES: t.CDefine = 256

# MemBuddy 伙伴系统常量
MEMBUDDY_MIN_BLOCK: t.CDefine = 32
MEMBUDDY_MAX_ORDERS: t.CDefine = 32
MEMBUDDY_HEADER_SIZE: t.CDefine = 8


def _align_up(val: t.CSizeT, align: t.CSizeT) -> t.CSizeT:
    if align == 0: return val
    return (val + align - 1) & ~(align - 1)


def _largest_pow2_le(val: t.CSizeT) -> t.CSizeT:
    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 = MEMBUDDY_MIN_BLOCK
    i: t.CInt
    for i in range(order):
        bs = bs << 1
    return bs


# ============================================================
# MemManager - 多态基类
# @t.CVTable 显式标记, 确保跨模块导入时也被识别为 vtable 类
# (HandlesImports 只通过装饰器检测 IsCVTable, 不像同模块自动推断)
# ============================================================
@t.CVTable
class MemManager:
    __provides__: list[str] = ['__memmgr__']
    base: t.CVoid | t.CPtr        # 内存区起始地址
    size: t.CSizeT                # 内存区总大小

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size

    # === 虚函数 (子类覆盖) ===

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        # 默认: 不支持分配
        return None

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        # 默认: 无操作
        return 0

    def reset(self) -> t.CInt:
        # 默认: 无操作
        return 0

    # === 具体方法 (调用虚函数, 自动多态分派) ===

    def calloc(self, count: t.CSizeT, size: t.CSizeT) -> t.CVoid | t.CPtr:
        total: t.CSizeT = count * size
        ptr: t.CVoid | t.CPtr = self.alloc(total)
        if ptr is not 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 is None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        new_ptr: t.CVoid | t.CPtr = self.alloc(new_size)
        if new_ptr is None:
            return ptr
        # 基类无法追踪旧分配大小, 不复制数据; 子类可覆盖以保留数据
        self.free(ptr)
        return new_ptr

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

    def __exit__(self):
        self.reset()

    def alloc_buf(self, capacity: t.CSizeT) -> viperio.Buf | t.CPtr:
        buf: viperio.Buf | t.CPtr = self.alloc(viperio.Buf.__sizeof__())
        if buf is None:
            return None
        data_ptr: t.CVoid | t.CPtr = self.alloc(capacity)
        if data_ptr is None:
            return None
        buf.__before_init__()
        buf.__init__(t.CChar(t.CUInt64T(data_ptr), t.CPtr), capacity)
        return buf


# ============================================================
# MemPool - arena bump 分配器
# 不支持单个 free, reset 一次性回收全部
# ============================================================
class MemPool(MemManager):
    offset: t.CSizeT           # bump 游标
    high_water: t.CSizeT       # 峰值水位

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size
        self.offset = 0
        self.high_water = 0

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        if size == 0: return None
        aligned: t.CSizeT = _align_up(size, MEMHUB_ALIGN)
        if self.offset + aligned > self.size: return None
        ptr: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.base) + self.offset, t.CPtr)
        self.offset += aligned
        if self.offset > self.high_water:
            self.high_water = self.offset
        return ptr

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        # bump 分配器不支持单个 free
        return 1

    def reset(self) -> t.CInt:
        self.offset = 0
        self.high_water = 0
        return 1


# ============================================================
# MemSlab - 定长块池
# arena 开头存位图, 后续区域按 block_size 切分, 空闲链管理
# ============================================================
class MemSlab(MemManager):
    block_size: t.CSizeT               # 每块大小 (对齐后)
    block_count: t.CSizeT              # 总块数
    used_count: t.CSizeT               # 已用块数
    free_list: t.CVoid | t.CPtr        # 空闲链头
    alloc_map: t.CUInt8T | t.CPtr      # 分配位图 (位于 arena 开头)
    alloc_map_size: t.CSizeT           # 位图字节数
    usable: t.CVoid | t.CPtr           # 可用块区起始 (跳过位图)
    usable_size: t.CSizeT              # 可用块区大小

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT, block_size: t.CSizeT):
        self.base = base
        self.size = size
        self.block_size = 0
        self.block_count = 0
        self.used_count = 0
        self.free_list = None
        self.alloc_map = None
        self.alloc_map_size = 0
        self.usable = None
        self.usable_size = 0

        if base is None: return

        bs: t.CSizeT = _align_up(block_size, MEMHUB_ALIGN)
        if bs < MEMSLAB_MIN_BLOCK: bs = MEMSLAB_MIN_BLOCK

        map_bytes: t.CSizeT = _align_up(MEMSLAB_BITMAP_BYTES, MEMHUB_ALIGN)
        if size < map_bytes + bs: return

        self.alloc_map = base
        self.alloc_map_size = map_bytes
        self.usable = t.CVoid(t.CUInt64T(base) + map_bytes, t.CPtr)
        self.usable_size = size - map_bytes
        self.block_size = bs

        # 清零位图
        idx: t.CSizeT = 0
        while idx < map_bytes:
            self.alloc_map[idx] = 0
            idx += 1

        self.block_count = self.usable_size / bs
        if self.block_count == 0: return

        # 构建空闲链: 每块首 8 字节存下一块指针
        self.free_list = None
        i: t.CSizeT = 0
        while i < self.block_count:
            block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.usable) + i * bs, t.CPtr)
            c.DerefAs(block, self.free_list)
            self.free_list = block
            i += 1

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        # slab 模式: 仅当请求 <= block_size 时分配一块
        if self.free_list is None: return None
        if size > self.block_size: return None
        block: t.CVoid | t.CPtr = self.free_list
        self.free_list = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        self.used_count += 1
        idx: t.CSizeT = (t.CUInt64T(block) - t.CUInt64T(self.usable)) / self.block_size
        self.alloc_map[idx / 8] = self.alloc_map[idx / 8] | t.CUInt8T(1 << (idx % 8))
        return block

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        if ptr is None: return 0
        p: t.CUInt64T = t.CUInt64T(ptr)
        if p < t.CUInt64T(self.usable): return 0
        if p >= t.CUInt64T(self.usable) + self.block_count * self.block_size: return 0
        if (p - t.CUInt64T(self.usable)) % self.block_size != 0: return 0
        idx: t.CSizeT = (p - t.CUInt64T(self.usable)) / self.block_size
        byte_idx: t.CSizeT = idx / 8
        bit_idx: t.CInt = t.CInt(idx % 8)
        if (self.alloc_map[byte_idx] >> bit_idx) & 1 == 0: return 0
        self.alloc_map[byte_idx] = self.alloc_map[byte_idx] & t.CUInt8T(~(1 << bit_idx))
        c.DerefAs(ptr, self.free_list)
        self.free_list = ptr
        self.used_count -= 1
        return 1

    def reset(self) -> t.CInt:
        self.used_count = 0
        self.free_list = None
        i: t.CSizeT = 0
        while i < self.block_count:
            block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.usable) + i * self.block_size, t.CPtr)
            c.DerefAs(block, self.free_list)
            self.free_list = block
            if i / 8 < self.alloc_map_size:
                self.alloc_map[i / 8] = 0
            i += 1
        return 1


# ============================================================
# MemBuddy - 伙伴系统分配器
# arena 开头存空闲链头数组, 后续区域按 2 的幂管理
# 线程安全 (自旋锁), 支持合并
# ============================================================
class MemBuddy(MemManager):
    max_order: t.CInt                  # 最大阶数
    free_lists: t.CUInt64T | t.CPtr    # 空闲链头数组 (位于 arena 开头)
    lock_val: t.CVolatile | t.CInt     # 自旋锁
    usable: t.CVoid | t.CPtr           # 可用区起始 (跳过 free_lists)
    usable_size: t.CSizeT              # 可用区大小 (2 的幂)

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size
        self.max_order = 0
        self.free_lists = None
        self.lock_val = 0
        self.usable = None
        self.usable_size = 0

        self.lock_val = 0
        fl_bytes: t.CSizeT = (MEMBUDDY_MAX_ORDERS + 1) * 8
        self.free_lists = base

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

        if size <= fl_bytes:
            self.usable = None
            self.usable_size = 0
            self.max_order = 0
            return

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

        if usable < MEMBUDDY_MIN_BLOCK:
            self.usable = None
            self.usable_size = 0
            self.max_order = 0
            return

        self.usable = t.CVoid(t.CUInt64T(base) + fl_bytes, t.CPtr)
        self.usable_size = usable

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

        # 整个可用区作为 max_order 阶空闲块
        self._fl_push(self.max_order, self.usable)

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

    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:
        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 is not 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.usable)
        bs: t.CSizeT = _block_size_at_order(order)
        buddy_offset: t.CSizeT = offset ^ bs
        return t.CVoid(t.CUInt64T(self.usable) + buddy_offset, t.CPtr)

    def _order_for_size(self, size: t.CSizeT) -> t.CInt:
        order: t.CInt = 0
        bs: t.CSizeT = MEMBUDDY_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:
        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)
        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:
        if ptr is None: return 0
        if self.usable is None: return 0
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_HEADER_SIZE, t.CPtr)
        if t.CUInt64T(block) < t.CUInt64T(self.usable): return 0
        if t.CUInt64T(block) >= t.CUInt64T(self.usable) + self.usable_size: return 0
        offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.usable)
        if offset % MEMBUDDY_MIN_BLOCK != 0: return 0
        stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        stored_val: t.CUInt64T = t.CUInt64T(stored)
        if (stored_val & 1) == 0: return 0
        order: t.CInt = t.CInt(stored_val >> 1)
        if order < 0: return 0
        if order > self.max_order: return 0
        return 1

    # === 锁 ===

    def _lock(self):
        while atom.__atomic_test_and_set(c.Addr(self.lock_val), atom.ATOMIC_ACQUIRE):
            pass

    def _unlock(self):
        atom.__atomic_clear(c.Addr(self.lock_val), atom.ATOMIC_RELEASE)

    # === 虚函数覆盖 ===

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        self._lock()
        result: t.CVoid | t.CPtr = None
        if self.usable is not None:
            if size != 0:
                needed: t.CSizeT = size + MEMBUDDY_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 is not None:
                        c.DerefAs(block, t.CVoid(t.CUInt64T((order << 1) | 1), t.CPtr))
                        result = t.CVoid(t.CUInt64T(block) + MEMBUDDY_HEADER_SIZE, t.CPtr)
        self._unlock()
        return result

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        self._lock()
        if ptr is not None:
            if self._is_valid_ptr(ptr) != 0:
                block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_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)
                c.DerefAs(block, t.CVoid(0, t.CPtr))
                self._coalesce(block, order)
        self._unlock()
        return 1

    def reset(self) -> t.CInt:
        if self.usable is None:
            return 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            self.free_lists[i] = 0
        self._fl_push(self.max_order, self.usable)
        return 1

    # realloc 覆盖: 伙伴系统从块头读取旧阶数, 无需 old_size
    def realloc(self, ptr: t.CVoid | t.CPtr, new_size: t.CSizeT) -> t.CVoid | t.CPtr:
        if ptr is None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        if self._is_valid_ptr(ptr) == 0:
            return None
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_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 + MEMBUDDY_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 is None:
            return ptr
        old_block_size: t.CSizeT = _block_size_at_order(old_order)
        old_data_size: t.CSizeT = old_block_size - MEMBUDDY_HEADER_SIZE
        string.memcpy(new_ptr, ptr, old_data_size)
        self.free(ptr)
        return new_ptr

    # === 状态查询 ===

    @property
    def mem_size(self) -> t.CSizeT:
        # 可用区总大小（2 的幂），分配上限为 mem_size - MEMBUDDY_HEADER_SIZE
        return self.usable_size

    def stats(self) -> t.CSizeT:
        # 返回可用区总大小，用于状态统计
        return self.usable_size

    def free_count(self) -> t.CSizeT:
        # 统计所有阶数空闲链中的块总数
        count: t.CSizeT = 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            head_val: t.CUInt64T = self.free_lists[i]
            cur: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
            while cur is not None:
                count += 1
                cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return count

    def self_check(self) -> t.CInt:
        # 验证伙伴分配器内部一致性，返回 0=OK，非 0=损坏
        if self.usable is None:
            return 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            head_val: t.CUInt64T = self.free_lists[i]
            cur: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
            while cur is not None:
                # 检查指针在可用区内
                if t.CUInt64T(cur) < t.CUInt64T(self.usable):
                    return 1
                if t.CUInt64T(cur) >= t.CUInt64T(self.usable) + self.usable_size:
                    return 2
                # 检查对齐：order i 的块必须按 _block_size_at_order(i) 对齐
                offset: t.CSizeT = t.CUInt64T(cur) - t.CUInt64T(self.usable)
                bs: t.CSizeT = _block_size_at_order(i)
                if offset % bs != 0:
                    return 3
                cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return 0
